Initial commit

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+checkpoints/*.pth filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,18 @@
+FROM python:3.9-slim
+
+ENV TRANSFORMERS_CACHE=/data/.cache/transformers
+ENV HF_HOME=/data/.cache/huggingface
+ENV MPLCONFIGDIR=/data/.cache/matplotlib
+
+WORKDIR /code
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+RUN mkdir -p checkpoints
+
+EXPOSE 7860
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -0,0 +1,9 @@
+---
+title: Room Type Classifier
+emoji: 🏠
+colorFrom: blue
+colorTo: orange
+sdk: docker
+app_file: main.py
+---
+

app.py

@@ -0,0 +1,81 @@
+import os
+import torch
+import torch.nn as nn
+import yaml
+from torchvision import models, transforms
+from PIL import Image
+import gradio as gr
+from transformers import ConvNextV2ForImageClassification
+
+CHECKPOINT_PATH = "checkpoints/room_classifier_best.pth"
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class HFConvNeXtWrapper(nn.Module):
+    def __init__(self, model_name, num_labels):
+        super(HFConvNeXtWrapper, self).__init__()
+        self.model = ConvNextV2ForImageClassification.from_pretrained(
+            model_name, num_labels=num_labels, ignore_mismatched_sizes=True)
+    def forward(self, x):
+        return self.model(x).logits
+
+def get_model(model_name, num_classes):
+    if model_name.startswith("efficientnet"):
+        model = models.efficientnet_b0(weights=None) if "b0" in model_name else models.efficientnet_b3(weights=None)
+        num_ftrs = model.classifier[1].in_features
+        model.classifier[1] = nn.Linear(num_ftrs, num_classes)
+    elif "convnextv2" in model_name:
+        model = HFConvNeXtWrapper(model_name, num_labels=num_classes)
+    elif model_name == "vit_b_16":
+        model = models.vit_b_16(weights=None)
+        model.heads.head = nn.Linear(model.heads.head.in_features, num_classes)
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+    return model
+
+if not os.path.exists(CHECKPOINT_PATH):
+    raise FileNotFoundError(f"Checkpoint not found at {CHECKPOINT_PATH}")
+
+print(f"Loading model from {CHECKPOINT_PATH}...")
+checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
+model_name = checkpoint['model_name']
+num_classes = checkpoint.get('num_classes', 5)
+
+class_to_idx = checkpoint.get('class_to_idx', None)
+if class_to_idx:
+    idx_to_class = {v: k for k, v in class_to_idx.items()}
+else:
+    print("Warning: class_to_idx not found in checkpoint. Using default 5 classes.")
+    idx_to_class = {0: 'Bathroom', 1: 'Bedroom', 2: 'Dining', 3: 'Kitchen', 4: 'Living'}
+
+model = get_model(model_name, num_classes)
+model.load_state_dict(checkpoint['state_dict'])
+model.to(DEVICE)
+model.eval()
+
+inference_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+def predict(pil_image):
+    if pil_image is None: return None
+    pil_image = pil_image.convert("RGB")
+    tensor = inference_transform(pil_image).unsqueeze(0).to(DEVICE)
+    
+    with torch.no_grad():
+        logits = model(tensor)
+        probs = torch.softmax(logits, dim=1).squeeze()
+        
+    return {idx_to_class[i]: float(probs[i]) for i in range(len(probs))}
+
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload Room Image"),
+    outputs=gr.Label(num_top_classes=5, label="Predictions"),
+    title="Room Type Classifier 🏠",
+    description=f"Classifies images into: {', '.join(idx_to_class.values())}",
+)
+
+if __name__ == "__main__":
+    iface.launch()

checkpoints/room_efficientnet_b0_best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d5ef02cf69916538affef9db11123ae3eecdb2175478284ad5341bfed055ebe5
+size 16360826

cm_config.yaml

@@ -0,0 +1,10 @@
+data_params:
+  image_size: 224      
+
+model_params:
+  name: "efficientnet_b0" 
+  num_classes: 5
+
+output_params:
+  save_dir: "checkpoints"
+  checkpoint_name: "room_efficientnet_b0_best.pth" 

main.py

@@ -0,0 +1,147 @@
+import os
+import torch
+import torch.nn as nn
+import yaml
+from torchvision import models, transforms
+from PIL import Image
+import gradio as gr
+import base64
+import io
+import time
+import threading
+from typing import List, Dict, Union, Optional
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from transformers import ConvNextV2ForImageClassification
+
+CHECKPOINT_DIR = "checkpoints"
+CONFIG_PATH = "cm_config.yaml"
+
+MODELS = {}
+LABELS = {}
+
+class HFConvNeXtWrapper(nn.Module):
+    def __init__(self, model_name, num_labels):
+        super(HFConvNeXtWrapper, self).__init__()
+        self.model = ConvNextV2ForImageClassification.from_pretrained(
+            model_name, num_labels=num_labels, ignore_mismatched_sizes=True)
+    def forward(self, x):
+        return self.model(x).logits
+
+def get_model(model_name, num_classes):
+    if model_name.startswith("efficientnet"):
+        model = models.efficientnet_b0(weights=None) if "b0" in model_name else models.efficientnet_b3(weights=None)
+        num_ftrs = model.classifier[1].in_features
+        model.classifier[1] = nn.Linear(num_ftrs, num_classes)
+    elif "convnextv2" in model_name:
+        model = HFConvNeXtWrapper(model_name, num_labels=num_classes)
+    elif model_name == "vit_b_16":
+        model = models.vit_b_16(weights=None)
+        model.heads.head = nn.Linear(model.heads.head.in_features, num_classes)
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+    return model
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+if not os.path.exists(CHECKPOINT_DIR):
+    os.makedirs(CHECKPOINT_DIR)
+
+model_files = [f for f in os.listdir(CHECKPOINT_DIR) if f.endswith('.pth')]
+default_model_name = None
+
+print(f"--- Loading models from {CHECKPOINT_DIR} ---")
+for filename in model_files:
+    path = os.path.join(CHECKPOINT_DIR, filename)
+    try:
+        ckpt = torch.load(path, map_location=device)
+        m_name = ckpt.get('model_name', 'efficientnet_b0')
+        n_classes = ckpt.get('num_classes', 5)
+        
+        model = get_model(m_name, n_classes)
+        model.load_state_dict(ckpt['state_dict'])
+        model.to(device)
+        model.eval()
+        
+        display_name = filename.replace('.pth', '')
+        MODELS[display_name] = model
+        
+        if 'class_to_idx' in ckpt:
+            LABELS[display_name] = {v: k for k, v in ckpt['class_to_idx'].items()}
+        else:
+            LABELS[display_name] = {0:'Bat', 1:'Bed', 2:'Din', 3:'Kit', 4:'Liv'}
+            
+        if default_model_name is None: default_model_name = display_name
+        print(f"Loaded: {display_name}")
+        
+    except Exception as e:
+        print(f"Failed to load {filename}: {e}")
+
+if not MODELS:
+    print("WARNING: No models loaded. Using Dummy for build.")
+    default_model_name = "dummy"
+
+inference_transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+class Base64Image(BaseModel): 
+    image_data: str
+    model_name: Optional[str] = default_model_name
+
+def base64_to_pil(base64_str: str) -> Image.Image:
+    if "base64," in base64_str: base64_str = base64_str.split("base64,")[1]
+    return Image.open(io.BytesIO(base64.b64decode(base64_str)))
+
+def run_inference(pil_image, model_key):
+    if model_key not in MODELS:
+        raise ValueError("Model not found")
+        
+    model = MODELS[model_key]
+    idx_map = LABELS[model_key]
+    
+    img_tensor = inference_transform(pil_image.convert("RGB")).unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        logits = model(img_tensor)
+        probs = torch.softmax(logits, dim=1).squeeze().tolist()
+    
+    return {idx_map[i]: float(probs[i]) for i in range(len(probs))}
+
+app = FastAPI(title="Room Type Classifier API")
+app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"])
+
+@app.get("/")
+def home():
+    return {"message": "Room Classifier API is running", "models": list(MODELS.keys())}
+
+@app.post("/predict")
+def predict_api(payload: Base64Image):
+    m_name = payload.model_name if payload.model_name else default_model_name
+    try:
+        img = base64_to_pil(payload.image_data)
+        result = run_inference(img, m_name)
+        return {"model": m_name, "predictions": result}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+def predict_gradio(img, model_choice):
+    if img is None: return None
+    return run_inference(img, model_choice)
+
+if MODELS:
+    gradio_iface = gr.Interface(
+        fn=predict_gradio,
+        inputs=[
+            gr.Image(type="pil", label="Image"),
+            gr.Dropdown(choices=list(MODELS.keys()), value=default_model_name, label="Model")
+        ],
+        outputs=gr.Label(num_top_classes=5),
+        title="Room Type Classifier",
+        description="Detects: Bathroom, Bedroom, Dining, Kitchen, Living",
+        allow_flagging="never"
+    )
+    app = gr.mount_gradio_app(app, gradio_iface, path="/gradio")

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+fastapi
+uvicorn
+gradio
+PyYAML
+python-multipart
+pydantic
+transformers