Update app.py

app.py

@@ -24,13 +24,8 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # Use GPU
 model = attempt_load(model_path, device=device)  # Placeholder for model loading
 model.eval()  # Set the model to evaluation mode
 
-#def preprocess_image(image_path):
 def preprocess_image(image):
-    #img0 = cv2.imread(image_path)
-    print("in preprocess-0 image.shape:",image.size)
     img = letterbox(image, 640, stride=32, auto=True)[0]  # Resize and pad to 640x640
-    #img = letterbox(img0, 640, stride=32, auto=True)[0]  # Resize and pad to 640x640
-    print("in preprocess-1 img.shape:",img.shape)
     img = img.transpose(2, 0, 1)[::-1]  # Convert BGR to RGB, 
     img = np.ascontiguousarray(img)
     img = torch.from_numpy(img).to(device)
@@ -38,10 +33,8 @@ def preprocess_image(image):
     img /= 255.0  # 0 - 255 to 0.0 - 1.0
     if img.ndimension() == 3:
         img = img.unsqueeze(0)
-    print("in preprocess-2 img.shape:",img.shape)
 
     return img, image
-    #return img, img0
 
 def infer(model, img):
     with torch.no_grad():
@@ -62,13 +55,11 @@ def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
     coords[:, :4].clip_(min=0, max=img1_shape[0])  # clip boxes
     return coords
 
-#def postprocess(pred, img0_shape, img):
 def postprocess(pred, img0, img):    
     pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False)
     results = []
     for det in pred:  # detections per image
         if len(det):
-            #det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0_shape).round()
             det[:, :4] = scale_coords(img.shape[2:], det[:, :4], img0.shape).round()
             for *xyxy, conf, cls in reversed(det):
                 results.append((xyxy, conf, cls))
@@ -76,16 +67,15 @@ def postprocess(pred, img0, img):
 
 def detect_objects(image_path):
     dicom_image, dicom_meta = read_and_preprocess_dicom(image_path)
-    #img, img0 = preprocess_image(image_path)
     img, img0 = preprocess_image(dicom_image)
     pred = infer(model, img)
-    #results = postprocess(pred, img0.shape, img)
     results = postprocess(pred, dicom_image, img)
     return results, dicom_image
 
 def draw_bounding_boxes(img, results):
     for (x1, y1, x2, y2), conf, cls in results:
         x1, y1, x2, y2 = map(int, [x1, y1, x2, y2])
+        print("results in draw_bounding box:",[(x1, y1, x2, y2), conf, cls])
         cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2)
         cv2.putText(img, f'{model.names[int(cls)]} {conf:.2f}', (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
     return img
@@ -94,7 +84,6 @@ def show_preds_image(filepath):
     results, img0 = detect_objects(filepath)
     print("in show preds:",results)
     #results = detect_objects(filepath)
-    #img0 = cv2.imread(filepath)
     img_with_boxes = draw_bounding_boxes(img0, results)
     return cv2.cvtColor(img_with_boxes, cv2.COLOR_BGR2RGB)