Adding app files

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -0,0 +1,144 @@
+# app.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import torchvision.transforms as T
+from PIL import Image, ImageDraw, ImageFont
+import gradio as gr
+from ultralytics import YOLO
+from transformers import ResNetModel
+import cv2
+
+class FlakeLayerClassifier(nn.Module):
+    def __init__(self, num_materials, material_dim, num_classes=4, dropout_prob=0.1, freeze_cnn=False):
+        super().__init__()
+        self.cnn = ResNetModel.from_pretrained("microsoft/resnet-18")
+        if freeze_cnn:
+            for p in self.cnn.parameters():
+                p.requires_grad = False
+
+        img_feat_dim = self.cnn.config.hidden_sizes[-1]
+        self.material_embedding = nn.Embedding(num_materials, material_dim)
+        self.dropout = nn.Dropout(dropout_prob)
+
+        self.fc_img = nn.Sequential(
+            nn.Linear(img_feat_dim, img_feat_dim),
+            nn.ReLU(inplace=True),
+            self.dropout,
+            nn.Linear(img_feat_dim, num_classes)
+        )
+
+        combined_dim = img_feat_dim + material_dim
+        self.fc_comb = nn.Sequential(
+            nn.Linear(combined_dim, combined_dim),
+            nn.ReLU(inplace=True),
+            self.dropout,
+            nn.Linear(combined_dim, num_classes)
+        )
+
+    def forward(self, pixel_values, material=None):
+        outputs = self.cnn(pixel_values=pixel_values)
+        img_feats = outputs.pooler_output.view(outputs.pooler_output.size(0), -1)
+
+        if material is None:
+            return self.fc_img(img_feats)
+
+        mat_emb = self.material_embedding(material)
+        combined = torch.cat([img_feats, mat_emb], dim=1)
+        return self.fc_comb(combined)
+
+def calibration(source_img, target_img):
+        source_lab = cv2.cvtColor(source_img, cv2.COLOR_BGR2LAB)
+        target_lab = cv2.cvtColor(target_img, cv2.COLOR_BGR2LAB)
+
+        for i in range(3):
+            src_mean, src_std = cv2.meanStdDev(source_lab[:, :, i])
+            tgt_mean, tgt_std = cv2.meanStdDev(target_lab[:, :, i])
+
+            target_lab[:, :, i] = (
+                (target_lab[:, :, i] - tgt_mean) * (src_std / tgt_std) + src_mean
+            ).clip(0, 255)
+
+        corrected_img = cv2.cvtColor(target_lab, cv2.COLOR_LAB2BGR)
+        return corrected_img.astype(np.uint8)
+
+
+device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+# Load YOLO detector
+#yolo = YOLO("/home/sankalp/flake_classification/models/best.pt")
+#yolo = YOLO("/home/sankalp/yolo_flake_detection/yolo11n_synthetic_runs/exp1/weights/best.pt")
+yolo = YOLO("/home/sankalp/yolo_flake_detection/yolo_runs/yolo11l_flake_runs/weights/best.pt")
+yolo.conf = 0.5
+
+# Load classifier weights
+ckpt = torch.load(
+    "/home/sankalp/flake_classification/models/flake_classifier.pth",
+    map_location=device
+)
+num_classes   = len(ckpt["class_to_idx"])
+classifier    = FlakeLayerClassifier(
+    num_materials=num_classes,
+    material_dim=64,
+    num_classes=num_classes,
+    dropout_prob=0.1,
+    freeze_cnn=False
+).to(device)
+classifier.load_state_dict(ckpt["model_state_dict"])
+classifier.eval()
+
+# Image processing transforms
+clf_tf = T.Compose([
+    T.Resize((224, 224)),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
+
+try:
+    FONT = ImageFont.truetype("arial.ttf", 20)
+except IOError:
+    FONT = ImageFont.load_default()
+
+# Inference + drawing
+def detect_and_classify(image: Image.Image):
+    #image = calibration(
+       # np.array(Image.open("/home/sankalp/gradio_flake_app/quantum-flake-pipeline/template/image.png")),
+        #np.array(image.convert("RGB")),
+    #)
+    #image = Image.fromarray(image)
+    img_rgb = np.array(image.convert("RGB"))
+    img_bgr = img_rgb[:, :, ::-1]
+    results = yolo(img_bgr, device=str(device))
+    boxes  = results[0].boxes.xyxy.cpu().numpy()
+    scores = results[0].boxes.conf.cpu().numpy()
+
+    draw = ImageDraw.Draw(image)
+    for (x1, y1, x2, y2), conf in zip(boxes, scores):
+        crop = image.crop((x1, y1, x2, y2))
+        inp  = clf_tf(crop).unsqueeze(0).to(device)  # (1,C,H,W)
+
+        with torch.no_grad():
+            logits = classifier(pixel_values=inp)
+            pred   = logits.argmax(1).item()
+            prob   = F.softmax(logits, dim=1)[0, pred].item()
+
+        label = f"Layer {pred+1} ({prob:.2f})"
+        # draw
+        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
+        draw.text((x1, max(0, y1-18)), label, fill="red", font=FONT)
+
+    return image
+
+# Gradio UI
+demo = gr.Interface(
+    fn=detect_and_classify,
+    inputs=gr.Image(type="pil", label="Upload Flake Image"),
+    outputs=gr.Image(type="pil", label="Annotated Output"),
+    title="Flake Detection + Layer Classification",
+    description="Upload an image → YOLO finds flakes → ResNet-18 head classifies their layer.",
+)
+
+if __name__ == "__main__":
+    demo.launch(share=True)

requirements.txt

@@ -0,0 +1,91 @@
+aiofiles==24.1.0
+annotated-types==0.7.0
+anyio==4.9.0
+certifi==2025.6.15
+charset-normalizer==3.4.2
+click==8.2.1
+contourpy==1.3.2
+cycler==0.12.1
+fastapi==0.115.13
+ffmpy==0.6.0
+filelock==3.18.0
+fonttools==4.58.4
+fsspec==2025.5.1
+gradio==5.34.2
+gradio_client==1.10.3
+groovy==0.1.2
+h11==0.16.0
+hf-xet==1.1.5
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.33.1
+idna==3.10
+Jinja2==3.1.6
+kiwisolver==1.4.8
+markdown-it-py==3.0.0
+MarkupSafe==3.0.2
+matplotlib==3.10.3
+mdurl==0.1.2
+mpmath==1.3.0
+networkx==3.5
+numpy==2.3.1
+nvidia-cublas-cu12==12.6.4.1
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cudnn-cu12==9.5.1.17
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cufile-cu12==1.11.1.6
+nvidia-curand-cu12==10.3.7.77
+nvidia-cusolver-cu12==11.7.1.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cusparselt-cu12==0.6.3
+nvidia-nccl-cu12==2.26.2
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-nvtx-cu12==12.6.77
+opencv-python==4.11.0.86
+orjson==3.10.18
+packaging==25.0
+pandas==2.3.0
+pillow==11.2.1
+psutil==7.0.0
+py-cpuinfo==9.0.0
+pydantic==2.11.7
+pydantic_core==2.33.2
+pydub==0.25.1
+Pygments==2.19.2
+pyparsing==3.2.3
+python-dateutil==2.9.0.post0
+python-multipart==0.0.20
+pytz==2025.2
+PyYAML==6.0.2
+regex==2024.11.6
+requests==2.32.4
+rich==14.0.0
+ruff==0.12.0
+safehttpx==0.1.6
+safetensors==0.5.3
+scipy==1.16.0
+semantic-version==2.10.0
+shellingham==1.5.4
+six==1.17.0
+sniffio==1.3.1
+starlette==0.46.2
+sympy==1.14.0
+tokenizers==0.21.2
+tomlkit==0.13.3
+torch==2.7.1
+torchaudio==2.7.1+cpu
+torchvision==0.22.1
+tqdm==4.67.1
+transformers==4.52.4
+triton==3.3.1
+typer==0.16.0
+typing-inspection==0.4.1
+typing_extensions==4.14.0
+tzdata==2025.2
+ultralytics==8.3.159
+ultralytics-thop==2.0.14
+urllib3==2.5.0
+uvicorn==0.34.3
+websockets==15.0.1