Update app.py

app.py

@@ -1,63 +1,65 @@
 import os
-import shutil
 from pathlib import Path
-import pandas as pd
-import gradio as gr
-import torch
+from tqdm.notebook import tqdm
 import numpy as np
+import pandas as pd
 from PIL import Image as PILImage
-from detectron2 import model_zoo
+import torch
+import cv2
+import pickle
+import shutil
 from detectron2.config import get_cfg
+from detectron2 import model_zoo
 from detectron2.engine import DefaultPredictor
 from detectron2.data import DatasetCatalog, MetadataCatalog
 from detectron2.utils.visualizer import ColorMode, Visualizer
-from tqdm import tqdm
-import uuid
-import cv2
-import pickle
 from math import ceil
-from typing import List, Dict, Any
-from numpy import ndarray
+import uuid
+from flask import Flask, request, send_file
 
-def get_vinbigdata_dicts_test(imgdir: Path, test_meta: pd.DataFrame, use_cache: bool = True, debug: bool = False):
+app = Flask(__name__)
+
+def get_vinbigdata_dicts_test(imgdir: Path, test_meta: pd.DataFrame, use_cache: bool = True, debug: bool = True):
     debug_str = f"_debug{int(debug)}"
     cache_path = Path(".") / f"dataset_dicts_cache_test{debug_str}.pkl"
     if not use_cache or not cache_path.exists():
+        print("Creating data...")
         if debug:
             test_meta = test_meta.iloc[:500]  # For debug
 
-        # Load 1 image to get image size
+        # Load 1 image to get image size.
         image_id = test_meta.loc[0, "image_id"]
-        image_path = imgdir / f"{image_id}.png"
-        image = cv2.imread(str(image_path))
+        image_path = os.path.join(imgdir, f"{image_id}.png")
+        image = cv2.imread(image_path)
         resized_height, resized_width, ch = image.shape
 
         dataset_dicts = []
         for index, test_meta_row in tqdm(test_meta.iterrows(), total=len(test_meta)):
             record = {}
             image_id, height, width = test_meta_row.values
-            filename = imgdir / f"{image_id}.png"
-            record["file_name"] = str(filename)
+            filename = os.path.join(imgdir, f"{image_id}.png")
+            record["file_name"] = filename
             record["image_id"] = image_id
             record["height"] = resized_height
             record["width"] = resized_width
             dataset_dicts.append(record)
-
+        
         with open(cache_path, mode="wb") as f:
             pickle.dump(dataset_dicts, f)
 
+    print(f"Load from cache {cache_path}")
     with open(cache_path, mode="rb") as f:
         dataset_dicts = pickle.load(f)
     return dataset_dicts
 
-def format_pred(labels: ndarray, boxes: ndarray, scores: ndarray) -> str:
+def format_pred(labels: np.ndarray, boxes: np.ndarray, scores: np.ndarray) -> str:
     pred_strings = []
     for label, score, bbox in zip(labels, scores, boxes):
         xmin, ymin, xmax, ymax = bbox.astype(np.int64)
         pred_strings.append(f"{label} {score} {xmin} {ymin} {xmax} {ymax}")
     return " ".join(pred_strings)
 
-def predict_batch(predictor: DefaultPredictor, im_list: List[ndarray]) -> List:
+def predict_batch(predictor: DefaultPredictor, im_list: list) -> list:
     with torch.no_grad():
         inputs_list = []
         for original_image in im_list:
@@ -92,15 +94,18 @@ def csv_create(new_image_path, image_id):
     test_meta_path = os.path.join(directory, 'test_meta.csv')
     test_meta_df.to_csv(test_meta_path, index=False)
 
+    print("CSV files have been generated successfully.")
     return sample_submission_path, test_meta_path
 
 def prediction(image_id_main, local_image_path, model_path):
     thing_classes = [
-        "Aortic enlargement", "Atelectasis", "Calcification", "Cardiomegaly", "Consolidation", 
-        "ILD", "Infiltration", "Lung Opacity", "Nodule/Mass", "Other lesion", "Pleural effusion", 
-        "Pleural thickening", "Pneumothorax", "Pulmonary fibrosis"
+        "Aortic enlargement", "Atelectasis", "Calcification", "Cardiomegaly",
+        "Consolidation", "ILD", "Infiltration", "Lung Opacity", "Nodule/Mass",
+        "Other lesion", "Pleural effusion", "Pleural thickening", "Pneumothorax", "Pulmonary fibrosis"
     ]
+    category_name_to_id = {class_name: index for index, class_name in enumerate(thing_classes)}
 
+    debug = False
     outdir = 'result_images'
     os.makedirs(outdir, exist_ok=True)
 
@@ -108,7 +113,7 @@ def prediction(image_id_main, local_image_path, model_path):
     os.makedirs(imgdir, exist_ok=True)
     shutil.copy(local_image_path, imgdir)
     new_image_path = os.path.join(imgdir, os.path.basename(local_image_path))
-    
+
     sample_submission, test_meta = csv_create(new_image_path, image_id_main)
 
     cfg = get_cfg()
@@ -116,24 +121,24 @@ def prediction(image_id_main, local_image_path, model_path):
     cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
     cfg.DATASETS.TEST = ()
     cfg.DATALOADER.NUM_WORKERS = 2
-    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml")
+    cfg.MODEL.WEIGHTS = model_path
     cfg.SOLVER.IMS_PER_BATCH = 2
     cfg.SOLVER.BASE_LR = 0.001
     cfg.SOLVER.MAX_ITER = 30000
     cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 512
     cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(thing_classes)
-
     cfg.MODEL.WEIGHTS = model_path
     cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.0
-    predictor = DefaultPredictor(cfg)
 
+    predictor = DefaultPredictor(cfg)
     unique_id = f"bigdata2_{uuid.uuid4().hex[:8]}"
-    DatasetCatalog.register(
-        unique_id, lambda: get_vinbigdata_dicts_test(imgdir, pd.read_csv(test_meta), debug=False)
-    )
+    DatasetCatalog.register(unique_id, lambda: get_vinbigdata_dicts_test(imgdir, pd.read_csv(test_meta), debug=debug))
     MetadataCatalog.get(unique_id).set(thing_classes=thing_classes)
     metadata = MetadataCatalog.get(unique_id)
-    dataset_dicts = get_vinbigdata_dicts_test(imgdir, pd.read_csv(test_meta), debug=False)
+    dataset_dicts = get_vinbigdata_dicts_test(imgdir, pd.read_csv(test_meta), debug=debug)
+
+    if debug:
+        dataset_dicts = dataset_dicts[:100]
 
     results_list = []
     batch_size = 4
@@ -145,19 +150,16 @@ def prediction(image_id_main, local_image_path, model_path):
         outputs_list = predict_batch(predictor, im_list)
 
         for im, outputs, d in zip(im_list, outputs_list, dataset_dicts_batch):
-            image_id, dim0, dim1 = pd.read_csv(test_meta).iloc[i].values
+            resized_height, resized_width, ch = im.shape
 
-            instances = outputs["instances"]
-            if len(instances) == 0:
-                result = {"image_id": image_id, "PredictionString": "14 1.0 0 0 1 1"}
-            else:
-                fields: Dict[str, Any] = instances.get_fields()
+            if outputs["instances"].has("pred_classes"):
+                fields = outputs["instances"].get_fields()
                 pred_classes = fields["pred_classes"]
                 pred_scores = fields["scores"]
                 pred_boxes = fields["pred_boxes"].tensor
 
-                h_ratio = dim0 / im.shape[0]
-                w_ratio = dim1 / im.shape[1]
+                h_ratio = d["height"] / resized_height
+                w_ratio = d["width"] / resized_width
                 pred_boxes[:, [0, 2]] *= w_ratio
                 pred_boxes[:, [1, 3]] *= h_ratio
 
@@ -166,11 +168,12 @@ def prediction(image_id_main, local_image_path, model_path):
                 pred_scores_array = pred_scores.cpu().numpy()
 
                 result = {
-                    "image_id": image_id,
-                    "PredictionString": format_pred(
-                        pred_classes_array, pred_boxes_array, pred_scores_array
-                    ),
+                    "image_id": d["image_id"],
+                    "PredictionString": format_pred(pred_classes_array, pred_boxes_array, pred_scores_array)
                 }
+            else:
+                result = {"image_id": d["image_id"], "PredictionString": "14 1 0 0 1 1"}
+            
             results_list.append(result)
 
     submission_det = pd.DataFrame(results_list, columns=['image_id', 'PredictionString'])
@@ -179,18 +182,19 @@ def prediction(image_id_main, local_image_path, model_path):
 
     return submission_det_path
 
-def process_image(image, image_id, model_path='path_to_your_model.pth'):
-    local_image_path = f'./{image_id}.png'
-    image.save(local_image_path)
+@app.route('/predict', methods=['POST'])
+def predict():
+    image_id = request.form['image_id']
+    image_file = request.files['image']
+    model_path = request.form['model_path']
+
+    local_image_path = os.path.join("input_images", image_file.filename)
+    os.makedirs("input_images", exist_ok=True)
+    image_file.save(local_image_path)
+
     submission_det_path = prediction(image_id, local_image_path, model_path)
-    return submission_det_path
 
-iface = gr.Interface(
-    fn=process_image,
-    inputs=[gr.Image(type="pil"), gr.Textbox(label="Image ID")],
-    outputs=gr.File(label="submission_det.csv"),
-    title="ECG Image Processing",
-    description="Upload an image and get the resulting CSV file."
-)
+    return send_file(submission_det_path, as_attachment=True)
 
-iface.launch()
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=8000)