Spaces:
Running
Running
| # Description: This is the main file to run the Gradio interface for the object detection model. | |
| from ultralytics import YOLO | |
| from PIL import Image | |
| import gradio as gr | |
| from huggingface_hub import snapshot_download | |
| import os | |
| model_path = "best_int8_openvino_model" | |
| # Example paths for Gradio | |
| image_examples = [["DurianMangosteen1.jpg"], ["DurianMangosteen2.jpg"]] | |
| # Load the model | |
| def load_model(repo_id): | |
| download_dir = snapshot_download(repo_id) # download the model from the Hugging Face Hub | |
| print(download_dir) | |
| path = os.path.join(download_dir, "best_int8_openvino_model") # path to the model | |
| print(path) | |
| detection_model = YOLO(path, task='detect') # load the model | |
| return detection_model | |
| # Predict the image | |
| def predict(pilimg): | |
| source = pilimg | |
| # x = np.asarray(pilimg) | |
| # print(x.shape) | |
| result = detection_model.predict(source, conf=0.4, iou=0.6) # confidence threshold, intersection over union threshold | |
| #print("Result: ", result) | |
| if not result or len(result[0].boxes) == 0: # if no object detected | |
| gr.Warning("No object detected in the image!") | |
| else: | |
| img_bgr = result[0].plot() # plot the image | |
| out_pilimg = Image.fromarray(img_bgr[..., ::-1]) # RGB-order PIL image | |
| return out_pilimg | |
| REPO_ID = "ITI107-2024S2/8035531F" # The repo ID of the model | |
| detection_model = load_model(REPO_ID) | |
| title = "Detect Durian and Mangosteen (King and Queen of Fruits) In The Image" | |
| interface = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="pil", label="Input Image"), | |
| outputs=gr.Image(type="pil", label="Object Detected Image"), | |
| title=title, | |
| examples=image_examples, | |
| ) | |
| # Launch the interface | |
| interface.launch(share=True) | |