Update app.py

app.py

@@ -22,8 +22,6 @@ from io import BytesIO
 # Load the YOLO model
 from models.common import DetectMultiBackend
 
-
-
 weights_path = "./last.pt"
 device = torch.device("cpu")  # Correctly define the device
 model = DetectMultiBackend(weights_path, device=device)  # Load YOLOv5 model correctly
@@ -49,24 +47,23 @@ transform = transforms.Compose([
 
 OBJECT_NAMES = ['enemies']
 
-def detect_objects_in_image(image):
 
+def detect_objects_in_image(image):
     # Ensure image is a PIL Image
+    data = Image.fromarray(array)
     if isinstance(image, torch.Tensor):
         image = transforms.ToPILImage()(image)  # Convert tensor to PIL image
-    
-    if isinstance(image, Image.Image):
+
+    if isinstance(data, Image.Image):
         orig_w, orig_h = image.size  # PIL image size returns (width, height)
     else:
-        # raise TypeError(f"Expected a PIL Image but got ")
+        raise TypeError(f"Expected a PIL Image but got ")
         print(type(image))
         exit()
 
     # Apply transformation
     img_tensor = transform(image).unsqueeze(0)
 
-
-
     with torch.no_grad():
         pred = model(img_tensor)[0]
 
@@ -106,7 +103,8 @@ def detect_objects_in_image(image):
             if object_name in object_counts:
                 object_counts[object_name] += 1
             cv2.rectangle(img_array, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            cv2.putText(img_array, f"{object_name}: {scores[i]:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+            cv2.putText(img_array, f"{object_name}: {scores[i]:.2f}", (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
+                        (0, 255, 0), 2)
 
     # Generate and return graph instead of dictionary
     graph_image = generate_vehicle_count_graph(object_counts)
@@ -116,58 +114,61 @@ def detect_objects_in_image(image):
 
 # def generate_vehicle_count_graph(object_counts):
 #     color_palette = ['#4C9ACD', '#88B8A3', '#7F9C9C', '#D1A3B5', '#A1C6EA', '#FFB6C1', '#F0E68C', '#D3B0D8', '#F8A5D1', '#B8B8D1']
-    
+
 #     fig, ax = plt.subplots(figsize=(8, 5))
 #     labels = list(object_counts.keys())
 #     values = list(object_counts.values())
-    
+
 #     ax.bar(labels, values, color=color_palette[:len(labels)])
-    
+
 #     ax.set_xlabel("Vehicle Categories", fontsize=12, fontweight='bold')
 #     ax.set_ylabel("Number of Vehicles", fontsize=12, fontweight='bold')
 #     ax.set_title("Detected Vehicles in Image", fontsize=14, fontweight='bold')
-    
+
 #     plt.xticks(rotation=45, ha='right', fontsize=10)
 #     plt.yticks(fontsize=10)
-    
+
 #     plt.tight_layout()
 
 #     buf = BytesIO()
 #     plt.savefig(buf, format='png')
 #     buf.seek(0)
-    
+
 #     return Image.open(buf)
 
 def generate_vehicle_count_graph(object_counts):
-    color_palette = ['#4C9ACD', '#88B8A3', '#7F9C9C', '#D1A3B5', '#A1C6EA', '#FFB6C1', '#F0E68C', '#D3B0D8', '#F8A5D1', '#B8B8D1']
-    
+    color_palette = ['#4C9ACD', '#88B8A3', '#7F9C9C', '#D1A3B5', '#A1C6EA', '#FFB6C1', '#F0E68C', '#D3B0D8', '#F8A5D1',
+                     '#B8B8D1']
+
     fig, ax = plt.subplots(figsize=(8, 5))
     labels = list(object_counts.keys())
     values = list(object_counts.values())
-    
+
     ax.bar(labels, values, color=color_palette[:len(labels)])
     ax.set_xlabel("Vehicle Categories", fontsize=12, fontweight='bold')
     ax.set_ylabel("Number of Vehicles", fontsize=12, fontweight='bold')
     ax.set_title("Detected Vehicles in Image", fontsize=14, fontweight='bold')
-    
+
     plt.xticks(rotation=45, ha='right', fontsize=10)
     plt.yticks(fontsize=10)
     plt.tight_layout()
-    
+
     buf = BytesIO()
     plt.savefig(buf, format='png')
     buf.seek(0)
-    
+
     plt.close(fig)  # ✅ CLOSE THE FIGURE TO FREE MEMORY
-    
+
     return Image.open(buf)
-    
+
+
 def detect_objects_in_video(video_input):
     cap = cv2.VideoCapture(video_input)
     if not cap.isOpened():
         return "Error: Cannot open video file.", None  # Returning a second value (None) to match expected outputs
 
-    frame_width, frame_height, fps = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FPS))
+    frame_width, frame_height, fps = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(
+        cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FPS))
     temp_video_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
     out = cv2.VideoWriter(temp_video_output, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
 
@@ -192,34 +193,37 @@ def detect_objects_in_video(video_input):
 
     return temp_video_output, graph_image  # Return both expected outputs
 
+
 def greet(name):
     return "Hello " + name + "!!"
 
+
 demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 
 from urllib.request import urlretrieve
 
 # get image examples from github
-urlretrieve("https://github.com/SamDaaLamb/ValorantTracker/blob/main/clip2_-1450-_jpg.jpg?raw=true", "clip2_-1450-_jpg.jpg") # make sure to use "copy image address when copying image from Github"
-urlretrieve("https://github.com/SamDaaLamb/ValorantTracker/blob/main/clip2_-539-_jpg.jpg?raw=true", "clip2_-539-_jpg.jpg")
-examples = [ # need to manually delete cache everytime new examples are added
-    ["clip2_-1450-_jpg.jpg"], 
+urlretrieve("https://github.com/SamDaaLamb/ValorantTracker/blob/main/clip2_-1450-_jpg.jpg?raw=true",
+            "clip2_-1450-_jpg.jpg")  # make sure to use "copy image address when copying image from Github"
+urlretrieve("https://github.com/SamDaaLamb/ValorantTracker/blob/main/clip2_-539-_jpg.jpg?raw=true",
+            "clip2_-539-_jpg.jpg")
+examples = [  # need to manually delete cache everytime new examples are added
+    ["clip2_-1450-_jpg.jpg"],
     ["clip2_-539-_jpg.jpg"]]
 
-
 # define app features and run
 title = "SpecLab Demo"
 description = "<p style='text-align: center'>Gradio demo for an ASPP model architecture trained on the SpecLab dataset. To use it, simply add your image, or click one of the examples to load them. Since this demo is run on CPU only, please allow additional time for processing. </p>"
 article = "<p style='text-align: center'><a href='https://github.com/Nano1337/SpecLab'>Github Repo</a></p>"
 css = "#0 {object-fit: contain;} #1 {object-fit: contain;}"
-demo = gr.Interface(fn=detect_objects_in_image, 
-                    title=title, 
+demo = gr.Interface(fn=detect_objects_in_image,
+                    title=title,
                     description=description,
                     article=article,
-                    inputs=gr.Image(elem_id=0, show_label=False), 
+                    inputs=gr.Image(elem_id=0, show_label=False),
                     outputs=gr.Image(elem_id=1, show_label=False),
-                    css=css, 
-                    examples=examples, 
+                    css=css,
+                    examples=examples,
                     cache_examples=True,
                     allow_flagging='never')
 demo.launch()