Spaces:
Sleeping
Sleeping
| from ultralytics import YOLO | |
| from PIL import Image | |
| import gradio as gr | |
| from huggingface_hub import snapshot_download | |
| import os | |
| import cv2 | |
| import tempfile | |
| from tqdm import tqdm | |
| def load_model(repo_id): | |
| download_dir = snapshot_download(repo_id) | |
| print(download_dir) | |
| path = os.path.join(download_dir, 'best_int8_openvino_model') | |
| print(path) | |
| detection_model = YOLO(path, task='detect') | |
| return detection_model | |
| def detect_image(pilimg): | |
| source = pilimg | |
| # x = np.asarray(pilimg) | |
| # print(x.shape) | |
| result = detection_model.predict(source, conf=0.5, iou=0.7) | |
| img_bgr = result[0].plot() | |
| out_pilimg = Image.fromarray(img_bgr[..., ::-1]) # RGB-order PIL image | |
| return out_pilimg | |
| def detect_video(video_file): | |
| video_reader = cv2.VideoCapture(video_file) | |
| nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT)) | |
| frame_width = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| frame_height = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| fps = video_reader.get(cv2.CAP_PROP_FPS) | |
| temp_video = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4') | |
| temp_video_path = temp_video.name | |
| video_writer = cv2.VideoWriter(temp_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps/2, (frame_width, frame_height)) | |
| for i in tqdm(range(nb_frames)): | |
| success, frame = video_reader.read() | |
| if success: | |
| results = detection_model.predict(frame, conf=0.5, iou=0.7) | |
| annotated_frame = results[0].plot() | |
| video_writer.write(annotated_frame) | |
| video_reader.release() | |
| video_writer.release() | |
| return temp_video_path | |
| REPO_ID = 'Xphosis/PotatoOrSweetPotato' | |
| detection_model = load_model(REPO_ID) | |
| with gr.Blocks() as app: | |
| with gr.Row(): | |
| with gr.Column(scale=1, min_width=300): | |
| # gr.Markdown('Upload image to detect Potato or Sweet Potato') | |
| img_input = gr.Image(type='pil', label='Upload Image to detect Potato or Sweet Potato',height=300,min_width=300) | |
| img_output = gr.Image(type='pil', label='Download Image',height=300,min_width=300) | |
| img_button = gr.Button('Analyze Image') | |
| img_button.click(detect_image, inputs=img_input, outputs=img_output,) | |
| clear_image_button = gr.Button('Clear Image') | |
| clear_image_button.click(fn=lambda: None, inputs=None, outputs=img_input) | |
| with gr.Column(scale=1, min_width=300): | |
| # gr.Markdown('Upload video to detect Potato or Sweet Potato') | |
| video_input = gr.Video(label='Upload Video to detect Potato or Sweet Potato', height=300,min_width=300) | |
| video_output = gr.Video(label='Download Video', autoplay=True, loop=True,show_share_button=True, show_download_button=True,height=300, min_width=300) | |
| video_button = gr.Button('Analyse Video') | |
| video_button.click(detect_video, inputs=video_input, outputs=video_output) | |
| clear_video_button = gr.Button('Clear Video') | |
| clear_video_button.click(fn=lambda: None, inputs=None, outputs=video_input) | |
| app.launch() |