0llheaven/CON-DETR-V5

Model Overview

This model is a DETR-based object detection model trained for medical image analysis, 4 classes : 0:Pneumonia, 1:Normal, 2:Pneumonia_bacteria, and 3:Pneumonia_virus

Model Description

Model Architecture: DEtection TRansformers (DETR) Training Data: Trained on a custom dataset of annotated medical images Intended Use: Designed for analyzing chest X-ray images to detect the presence and type of pneumonia, or classify as normal..

Uses

Example Code test model!

import os 
from transformers import AutoImageProcessor, AutoModelForObjectDetection
import torch
from PIL import Image
import pandas as pd

folder_path = ""
processor = AutoImageProcessor.from_pretrained("0llheaven/CON-DETR-V5")
model = AutoModelForObjectDetection.from_pretrained("0llheaven/CON-DETR-V5")

results_list = []

for image_name in os.listdir(folder_path):
    if image_name.endswith((".jpg", ".png", ".jpeg")): 
        image_path = os.path.join(folder_path, image_name)
        image = Image.open(image_path)

        # RGB To grayscale
        if image.mode != "RGB":
            image = image.convert("RGB")

        # prediction
        inputs = processor(images=image, return_tensors="pt")
        outputs = model(**inputs)

        target_sizes = torch.tensor([image.size[::-1]])
        results = processor.post_process_object_detection(outputs, target_sizes=target_sizes)
        print(f"Processing image: {image_name}")

        # detect_box
        detected_any = False
        for result in results:
            scores = result["scores"]
            labels = result["labels"]
            boxes = result["boxes"]

            filtered_data = [(score, label, box) for score, label, box in zip(scores, labels, boxes) if score > 0.5][:2] # จำกัดให้เก็บได้ไม่เกิน 2 กล่อง
            for score, label, box in zip(scores, labels, boxes):
                if score > 0.5:
                    if len(filtered_data) > 0:
                        detected_any = True
                            
                    for score, label, box in filtered_data:
                        if label.item() == 0:
                            label_name = "Pneumonia"
                        elif label.item() == 1:
                            label_name = "Normal"
                        elif label.item() == 2:
                            label_name = "Pneumonia_bacteria"
                        else:
                            label_name = "Pneumonia_virus"
                        
                        xmin, ymin, xmax, ymax = [round(i, 2) for i in box.tolist()]
                        print(f" - Detected {label_name} with score {round(score.item(), 3)} at {xmin, ymin, xmax, ymax}")
                        results_list.append({
                            "image_name": image_name,
                            "label": label_name,
                            "xmin": xmin,
                            "ymin": ymin,
                            "xmax": xmax,
                            "ymax": ymax,
                            "score": round(score.item(), 3),
                        })

        if not detected_any:
            print(" - No Detect")
            results_list.append({
                "image_name": image_name,
                "label": "Other",
                "xmin": 0,
                "ymin": 0,
                "xmax": 0,
                "ymax": 0,
                "score": 0,
            })

results_df = pd.DataFrame(results_list)
print("\nFinal results:")
results_df.to_csv("testmodel.csv", index=False)