Update app.py

app.py

@@ -1,24 +1,380 @@
 # app.py
+import os
+import io
+import zipfile
+import json
+import shutil
+from pathlib import Path
+from PIL import Image
+import torch
+from torchvision import models, transforms, datasets
+from torch.utils.data import DataLoader, random_split
+import torch.nn as nn
+import torch.optim as optim
 import gradio as gr
-from transformers import pipeline
+import time
 
-# Load a ready-made image classification model
-classifier = pipeline("image-classification", model="google/vit-base-patch16-224")
+ROOT = Path(".")
+DATA_ZIP_NAME = "dataset.zip"   # upload your Roboflow export here
+WORK_DIR = ROOT / "roboflow_dataset"
+CLASSIFY_DIR = ROOT / "classification_data"
+MODEL_PATH = ROOT / "model.pth"
+CLASSES_JSON = ROOT / "classes.json"
 
-# Define prediction function
-def predict_image(img):
-    results = classifier(img)
-    # Take top 3 predictions
-    return {res["label"]: float(res["score"]) for res in results[:3]}
+# Training config (tweak if needed)
+BATCH_SIZE = 16
+IMG_SIZE = 224
+NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", 3))   # small default for Spaces CPU
+LR = 1e-3
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Create Gradio interface
+
+def safe_mkdir(p: Path):
+    p.mkdir(parents=True, exist_ok=True)
+
+
+def extract_zip_to_workdir(zip_path: Path, out_dir: Path):
+    if out_dir.exists():
+        shutil.rmtree(out_dir)
+    safe_mkdir(out_dir)
+    with zipfile.ZipFile(zip_path, "r") as z:
+        z.extractall(out_dir)
+
+
+def find_classes_mapping(workdir: Path):
+    # Roboflow usually includes a data.yaml or classes.txt or a names list.
+    # Try common locations.
+    data_yaml = workdir / "data.yaml"
+    classes_txt = workdir / "classes.txt"
+    # Sometimes Roboflow includes a folder "labels" and a file "labels.names" or "classes.txt"
+    if classes_txt.exists():
+        names = [x.strip() for x in classes_txt.read_text().splitlines() if x.strip()]
+        return names
+    if data_yaml.exists():
+        import yaml  # note: pyyaml must be in requirements if needed
+        try:
+            parsed = yaml.safe_load(data_yaml.read_text())
+            if "names" in parsed:
+                # could be list or dict
+                n = parsed["names"]
+                if isinstance(n, dict):
+                    return [n[k] for k in sorted(n.keys(), key=lambda x: int(x))]
+                elif isinstance(n, list):
+                    return n
+        except Exception:
+            pass
+    # fallback: try to find a file named "classes.txt" or "labels.names"
+    for candidate in workdir.rglob("classes.txt"):
+        names = [x.strip() for x in candidate.read_text().splitlines() if x.strip()]
+        if names:
+            return names
+    for candidate in workdir.rglob("labels.names"):
+        names = [x.strip() for x in candidate.read_text().splitlines() if x.strip()]
+        if names:
+            return names
+    # last resort: scan label files to get max class index, produce numeric names
+    max_idx = -1
+    for lbl in workdir.rglob("labels/*.txt"):
+        for line in lbl.read_text().splitlines():
+            parts = line.strip().split()
+            if len(parts) >= 1:
+                try:
+                    idx = int(float(parts[0]))
+                    max_idx = max(max_idx, idx)
+                except:
+                    pass
+    if max_idx >= 0:
+        return [f"class_{i}" for i in range(max_idx + 1)]
+    return []
+
+
+def convert_roboflow_detection_to_classification(workdir: Path, outdir: Path):
+    """
+    Creates a folder-structured classification dataset:
+    outdir/train/<class_name>/*.jpg
+    outdir/valid/<class_name>/*.jpg
+    It uses label files (YOLO txt) to assign the main class for each image.
+    If bounding box info is available, it crops the bbox; otherwise it copies the image.
+    """
+    if outdir.exists():
+        shutil.rmtree(outdir)
+    safe_mkdir(outdir)
+
+    # Try common image and label folders
+    images_dirs = []
+    labels_dirs = []
+    for p in workdir.iterdir():
+        if p.is_dir():
+            if p.name.lower() in ("images", "image", "images/train", "train", "valid", "test"):
+                images_dirs.append(p)
+            if p.name.lower() in ("labels", "annotations"):
+                labels_dirs.append(p)
+
+    # simpler approach: look for 'images' and 'labels' in any depth
+    images_all = list(workdir.rglob("images/*")) + list(workdir.rglob("images/*/*"))
+    if not images_all:
+        # fallback to all popular image file types in workdir
+        images_all = [p for p in workdir.rglob("*") if p.suffix.lower() in (".jpg", ".jpeg", ".png")]
+
+    # mapping of image filename (no path) to its full path
+    img_map = {p.name: p for p in images_all}
+
+    # find label files
+    label_files = list(workdir.rglob("labels/*.txt")) + list(workdir.rglob("labels/*/*.txt"))
+    if not label_files:
+        # some exports put labels alongside images with same base name but different extension
+        label_files = [p for p in workdir.rglob("*.txt") if p.stem in img_map]
+
+    # find class names
+    classes = find_classes_mapping(workdir)
+    if not classes:
+        # if not available, default to single class "unknown"
+        classes = ["class_0"]
+
+    # prepare train/valid split target folders (Roboflow often has train/valid folders; try to preserve)
+    # We'll just create train and valid
+    train_out = outdir / "train"
+    valid_out = outdir / "valid"
+    safe_mkdir(train_out)
+    safe_mkdir(valid_out)
+
+    # Load label->image mapping from label_files
+    # We'll assume label files mirror the image names: e.g., images/train/img1.jpg and labels/train/img1.txt
+    img_to_labels = {}
+    for lbl in label_files:
+        name = lbl.stem
+        if name in img_map:
+            img_to_labels[name] = lbl
+
+    # If Roboflow has images split into train/valid dirs, detect them
+    # Otherwise we'll create a split based on filenames (80/20)
+    # Build a dataset list
+    dataset_rows = []
+    for img_name, img_path in img_map.items():
+        lbl = img_to_labels.get(Path(img_name).stem)
+        # Determine main class for this image (first label line)
+        main_class = None
+        bbox = None
+        if lbl and lbl.exists():
+            lines = [l for l in lbl.read_text().splitlines() if l.strip()]
+            if lines:
+                parts = lines[0].split()
+                try:
+                    cls_idx = int(float(parts[0]))
+                    main_class = classes[cls_idx] if cls_idx < len(classes) else f"class_{cls_idx}"
+                    if len(parts) >= 5:
+                        # YOLO format: cls x_center y_center width height (normalized)
+                        bbox = tuple(float(x) for x in parts[1:5])
+                except Exception:
+                    pass
+        if not main_class:
+            # fallback: mark as unknown
+            main_class = "unknown"
+            if "unknown" not in classes:
+                classes.append("unknown")
+        dataset_rows.append((img_path, main_class, bbox))
+
+    # do deterministic split
+    dataset_rows.sort(key=lambda x: x[0].name)
+    split_idx = int(0.8 * len(dataset_rows))
+    train_rows = dataset_rows[:split_idx]
+    valid_rows = dataset_rows[split_idx:]
+
+    def save_rows(rows, dest_folder):
+        for img_path, cls_name, bbox in rows:
+            dest_cls = dest_folder / cls_name
+            safe_mkdir(dest_cls)
+            try:
+                img = Image.open(img_path).convert("RGB")
+                if bbox:
+                    # bbox are normalized; convert to pixel coords
+                    w, h = img.size
+                    xc, yc, bw, bh = bbox
+                    left = int((xc - bw / 2) * w)
+                    right = int((xc + bw / 2) * w)
+                    top = int((yc - bh / 2) * h)
+                    bottom = int((yc + bh / 2) * h)
+                    # clamp
+                    left = max(0, left); right = min(w, right)
+                    top = max(0, top); bottom = min(h, bottom)
+                    if right - left > 10 and bottom - top > 10:
+                        img = img.crop((left, top, right, bottom))
+                # save with a unique name
+                dest_path = dest_cls / img_path.name
+                img.save(dest_path)
+            except Exception as e:
+                print("Skipping", img_path, "due to", e)
+
+    save_rows(train_rows, train_out)
+    save_rows(valid_rows, valid_out)
+
+    # Save classes json
+    with open(CLASSES_JSON, "w") as f:
+        json.dump(classes, f)
+    return classes
+
+
+def build_model(num_classes):
+    model = models.resnet18(pretrained=True)
+    in_features = model.fc.in_features
+    model.fc = nn.Linear(in_features, num_classes)
+    return model
+
+
+def train_model(data_dir: Path, classes):
+    print("Starting training. This may take some time on CPU.")
+    num_classes = len(classes)
+    model = build_model(num_classes).to(DEVICE)
+
+    transform_train = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.RandomHorizontalFlip(),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+    ])
+    transform_valid = transforms.Compose([
+        transforms.Resize((IMG_SIZE, IMG_SIZE)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+    ])
+
+    dataset_train = datasets.ImageFolder(str(data_dir / "train"), transform=transform_train)
+    dataset_valid = datasets.ImageFolder(str(data_dir / "valid"), transform=transform_valid)
+
+    # If ImageFolder class mapping differs from classes list, use folder names.
+    # Dataloaders
+    if len(dataset_train) == 0:
+        raise RuntimeError("No training images found. Please check dataset structure.")
+
+    loader_train = DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
+    loader_valid = DataLoader(dataset_valid, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
+
+    criterion = nn.CrossEntropyLoss()
+    optimizer = optim.Adam(model.parameters(), lr=LR)
+
+    best_val = 0.0
+    for epoch in range(NUM_EPOCHS):
+        model.train()
+        running = 0.0
+        for imgs, labels in loader_train:
+            imgs = imgs.to(DEVICE)
+            labels = labels.to(DEVICE)
+            optimizer.zero_grad()
+            outputs = model(imgs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+            running += loss.item()
+        # validation
+        model.eval()
+        correct = 0
+        total = 0
+        with torch.no_grad():
+            for imgs, labels in loader_valid:
+                imgs = imgs.to(DEVICE)
+                labels = labels.to(DEVICE)
+                outputs = model(imgs)
+                _, preds = torch.max(outputs, 1)
+                correct += (preds == labels).sum().item()
+                total += labels.size(0)
+        acc = correct / total if total > 0 else 0.0
+        print(f"Epoch {epoch+1}/{NUM_EPOCHS}, loss={running:.4f}, val_acc={acc:.4f}")
+        if acc > best_val:
+            best_val = acc
+            # save best
+            torch.save({
+                "model_state": model.state_dict(),
+                "classes": classes
+            }, MODEL_PATH)
+    print("Training complete. Best val acc:", best_val)
+    # final save if not saved
+    if not MODEL_PATH.exists():
+        torch.save({
+            "model_state": model.state_dict(),
+            "classes": classes
+        }, MODEL_PATH)
+    return MODEL_PATH.exists()
+
+
+def load_saved_model(path: Path):
+    data = torch.load(path, map_location=DEVICE)
+    classes = data.get("classes", None)
+    if not classes and Path(CLASSES_JSON).exists():
+        classes = json.loads(Path(CLASSES_JSON).read_text())
+    if not classes:
+        classes = [f"class_{i}" for i in range(2)]
+    model = build_model(len(classes))
+    model.load_state_dict(data["model_state"])
+    model.to(DEVICE).eval()
+    return model, classes
+
+
+# Prepare model at startup
+MODEL = None
+BREEDS = None
+
+def startup():
+    global MODEL, BREEDS
+    # If model exists, load directly
+    if Path(MODEL_PATH).exists():
+        try:
+            MODEL, BREEDS = load_saved_model(Path(MODEL_PATH))
+            print("Loaded existing model with classes:", BREEDS)
+            return
+        except Exception as e:
+            print("Failed to load existing model:", e)
+
+    # If dataset.zip exists, extract and convert, then train
+    if Path(DATA_ZIP_NAME).exists():
+        print("dataset.zip found. Extracting and preparing...")
+        extract_zip_to_workdir(Path(DATA_ZIP_NAME), WORK_DIR)
+        classes = convert_roboflow_detection_to_classification(WORK_DIR, CLASSIFY_DIR)
+        print("Prepared classification dataset with classes:", classes)
+        # train (may be slow on CPU)
+        try:
+            trained = train_model(CLASSIFY_DIR, classes)
+            if trained:
+                MODEL, BREEDS = load_saved_model(Path(MODEL_PATH))
+        except Exception as e:
+            print("Training failed:", e)
+    else:
+        print("No dataset.zip found. Please upload dataset.zip to the Space root or upload a model.pth")
+
+
+# Prediction function
+transform_predict = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
+])
+
+def predict_image(pil_img):
+    global MODEL, BREEDS
+    if MODEL is None:
+        return {"error": "Model not ready. Upload dataset.zip to train, or model.pth to load."}
+    img = pil_img.convert("RGB")
+    x = transform_predict(img).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        out = MODEL(x)
+        probs = torch.nn.functional.softmax(out[0], dim=0).cpu().numpy()
+    # top 3
+    indices = probs.argsort()[::-1][:3]
+    return {BREEDS[int(i)]: float(probs[int(i)]) for i in indices}
+
+# Run startup (this will attempt to load or train)
+start_time = time.time()
+startup()
+print("Startup complete in", time.time() - start_time, "seconds")
+
+# Build Gradio app
 demo = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(type="pil"),
     outputs=gr.Label(num_top_classes=3),
-    title="Simple Image Classifier",
-    description="Upload an image and let AI classify it!"
+    examples=[],
+    title="Cow Breed Classifier",
+    description="Upload a cow image. If you uploaded Roboflow dataset.zip to the Space root, the Space will auto-train on start (small number of epochs). If you already have a trained model.pth, upload that instead to skip training."
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)