Create app.py

app.py

@@ -0,0 +1,196 @@
+import os
+import shutil
+from pathlib import Path
+import pandas as pd
+import gradio as gr
+import torch
+import numpy as np
+from PIL import Image as PILImage
+from detectron2 import model_zoo
+from detectron2.config import get_cfg
+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
+
+def get_vinbigdata_dicts_test(imgdir: Path, test_meta: pd.DataFrame, use_cache: bool = True, debug: bool = False):
+    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():
+        if debug:
+            test_meta = test_meta.iloc[:500]  # For debug
+
+        # 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))
+        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)
+            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)
+
+    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:
+    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:
+    with torch.no_grad():
+        inputs_list = []
+        for original_image in im_list:
+            if predictor.input_format == "RGB":
+                original_image = original_image[:, :, ::-1]
+            height, width = original_image.shape[:2]
+            image = torch.as_tensor(original_image.astype("float32").transpose(2, 0, 1))
+            inputs = {"image": image, "height": height, "width": width}
+            inputs_list.append(inputs)
+        predictions = predictor.model(inputs_list)
+        return predictions
+
+def csv_create(new_image_path, image_id):
+    image = PILImage.open(new_image_path)
+    width, height = image.size
+    directory = os.path.dirname(new_image_path)
+
+    sample_submission_data = {
+        'image_id': [image_id],
+        'PredictionString': ['14 1 0 0 1 1']
+    }
+    sample_submission_df = pd.DataFrame(sample_submission_data)
+    sample_submission_path = os.path.join(directory, 'sample_submission.csv')
+    sample_submission_df.to_csv(sample_submission_path, index=False)
+
+    test_meta_data = {
+        'image_id': [image_id],
+        'dim0': [width],
+        'dim1': [height]
+    }
+    test_meta_df = pd.DataFrame(test_meta_data)
+    test_meta_path = os.path.join(directory, 'test_meta.csv')
+    test_meta_df.to_csv(test_meta_path, index=False)
+
+    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"
+    ]
+
+    outdir = 'result_images'
+    os.makedirs(outdir, exist_ok=True)
+
+    imgdir = f'processed_images_{image_id_main}'
+    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()
+    cfg.OUTPUT_DIR = outdir
+    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.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)
+
+    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)
+    )
+    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)
+
+    results_list = []
+    batch_size = 4
+
+    for i in tqdm(range(ceil(len(dataset_dicts) / batch_size))):
+        inds = list(range(batch_size * i, min(batch_size * (i + 1), len(dataset_dicts))))
+        dataset_dicts_batch = [dataset_dicts[i] for i in inds]
+        im_list = [cv2.imread(d["file_name"]) for d in dataset_dicts_batch]
+        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
+
+            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()
+                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]
+                pred_boxes[:, [0, 2]] *= w_ratio
+                pred_boxes[:, [1, 3]] *= h_ratio
+
+                pred_classes_array = pred_classes.cpu().numpy()
+                pred_boxes_array = pred_boxes.cpu().numpy()
+                pred_scores_array = pred_scores.cpu().numpy()
+
+                result = {
+                    "image_id": image_id,
+                    "PredictionString": format_pred(
+                        pred_classes_array, pred_boxes_array, pred_scores_array
+                    ),
+                }
+            results_list.append(result)
+
+    submission_det = pd.DataFrame(results_list, columns=['image_id', 'PredictionString'])
+    submission_det_path = os.path.join(outdir, "submission_det.csv")
+    submission_det.to_csv(submission_det_path, index=False)
+
+    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)
+    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."
+)
+
+iface.launch()