Trying to get YOLOV8

app.py

@@ -6,6 +6,7 @@ import torch
 import pathlib
 from PIL import Image
 from transformers import AutoFeatureExtractor, DetrForObjectDetection, YolosForObjectDetection
+from ultralyticsplus import YOLO, render_result
 
 import os
 
@@ -58,48 +59,71 @@ def detect_objects(model_name,url_input,image_input,threshold):
     
     #Extract model and feature extractor
     feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
-    
-    if 'detr' in model_name:
-        
-        model = DetrForObjectDetection.from_pretrained(model_name)
+
+    if 'yolov8' in model_name:
+
+        model = YOLO(model_name)
+        # set model parameters
+        model.overrides['conf'] = 0.25  # NMS confidence threshold
+        model.overrides['iou'] = 0.45  # NMS IoU threshold
+        model.overrides['agnostic_nms'] = False  # NMS class-agnostic
+        model.overrides['max_det'] = 1000  # maximum number of detections per image
+
+        results = model.predict(image_input)
+
+        render = render_result(model=model, image=image_input, result=results[0])
+
+        return render, ""
         
-    elif 'yolos' in model_name:
+        # for result in results:
+        #     # https://docs.ultralytics.com/modes/predict/#key-features-of-predict-mode 
+        #     #TODO
+        #     im_array = result.plot()
+        #     im = Image.fromarray(im_array[..., ::=1])
+            
     
-        model = YolosForObjectDetection.from_pretrained(model_name)
+    else:
+        if 'detr' in model_name:
+        
+            model = DetrForObjectDetection.from_pretrained(model_name)
 
-    tb_label = ""
-    if validators.url(url_input):
-        image = Image.open(requests.get(url_input, stream=True).raw)
-        tb_label = "Confidence Values URL"
+        elif 'yolos' in model_name:
         
-    elif image_input:
-        image = image_input
-        tb_label = "Confidence Values Upload"
+            model = YolosForObjectDetection.from_pretrained(model_name)
     
-    #Make prediction
-    processed_output_list = make_prediction(image, feature_extractor, model)
-    print("After make_prediction" + str(processed_output_list))
-    processed_outputs = processed_output_list[0]
+        tb_label = ""
+        if validators.url(url_input):
+            image = Image.open(requests.get(url_input, stream=True).raw)
+            tb_label = "Confidence Values URL"
+            
+        elif image_input:
+            image = image_input
+            tb_label = "Confidence Values Upload"
+        
+        #Make prediction
+        processed_output_list = make_prediction(image, feature_extractor, model)
+        print("After make_prediction" + str(processed_output_list))
+        processed_outputs = processed_output_list[0]
+        
+        #Visualize prediction
+        viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
+        
+        # return [viz_img, processed_outputs]
+        # print(type(viz_img))
     
-    #Visualize prediction
-    viz_img = visualize_prediction(image, processed_outputs, threshold, model.config.id2label)
+        final_str_abv = ""
+        final_str_else = ""
+        for score, label, box in sorted(zip(processed_outputs["scores"], processed_outputs["labels"], processed_outputs["boxes"]), key = lambda x: x[0].item(), reverse=True):
+            box = [round(i, 2) for i in box.tolist()]
+            if score.item() >= threshold:
+                final_str_abv += f"Detected `{model.config.id2label[label.item()]}` with confidence `{round(score.item(), 3)}` at location `{box}`\n"
+            else:
+                final_str_else += f"Detected `{model.config.id2label[label.item()]}` with confidence `{round(score.item(), 3)}` at location `{box}`\n"
     
-    # return [viz_img, processed_outputs]
-    # print(type(viz_img))
-
-    final_str_abv = ""
-    final_str_else = ""
-    for score, label, box in sorted(zip(processed_outputs["scores"], processed_outputs["labels"], processed_outputs["boxes"]), key = lambda x: x[0].item(), reverse=True):
-        box = [round(i, 2) for i in box.tolist()]
-        if score.item() >= threshold:
-            final_str_abv += f"Detected `{model.config.id2label[label.item()]}` with confidence `{round(score.item(), 3)}` at location `{box}`\n"
-        else:
-            final_str_else += f"Detected `{model.config.id2label[label.item()]}` with confidence `{round(score.item(), 3)}` at location `{box}`\n"
-
-    # https://docs.python.org/3/library/string.html#format-examples
-    final_str = "{:*^50}\n".format("ABOVE THRESHOLD OR EQUAL") + final_str_abv + "\n{:*^50}\n".format("BELOW THRESHOLD")+final_str_else
-        
-    return viz_img, final_str
+        # https://docs.python.org/3/library/string.html#format-examples
+        final_str = "{:*^50}\n".format("ABOVE THRESHOLD OR EQUAL") + final_str_abv + "\n{:*^50}\n".format("BELOW THRESHOLD")+final_str_else
+            
+        return viz_img, final_str
         
 def set_example_image(example: list) -> dict:
     return gr.Image.update(value=example[0])
@@ -119,10 +143,11 @@ Links to HuggingFace Models:
 - [hustvl/yolos-tiny](https://huggingface.co/hustvl/yolos-tiny)
 - [facebook/detr-resnet-101-dc5](https://huggingface.co/facebook/detr-resnet-101-dc5)
 - [hustvl/yolos-small-300](https://huggingface.co/hustvl/yolos-small-300)
+- [mshamrai/yolov8x-visdrone](https://huggingface.co/mshamrai/yolov8x-visdrone)
 
 """
 
-models = ["facebook/detr-resnet-50","facebook/detr-resnet-101",'hustvl/yolos-small','hustvl/yolos-tiny','facebook/detr-resnet-101-dc5', 'hustvl/yolos-small-300']
+models = ["facebook/detr-resnet-50","facebook/detr-resnet-101",'hustvl/yolos-small','hustvl/yolos-tiny','facebook/detr-resnet-101-dc5', 'hustvl/yolos-small-300', 'mshamrai/yolov8x-visdrone']
 urls = ["https://c8.alamy.com/comp/J2AB4K/the-new-york-stock-exchange-on-the-wall-street-in-new-york-J2AB4K.jpg"]
 
 # twitter_link = """