app.py

#!/usr/bin/env python3
import os

# ────────────────────────────────────────────────────────────
# 1) FORCE HOME into /tmp so that ~/.streamlit is writable
# ────────────────────────────────────────────────────────────
os.environ["HOME"] = "/tmp"
streamlit_config = os.path.join(os.environ["HOME"], ".streamlit")
os.makedirs(streamlit_config, exist_ok=True)
os.environ["STREAMLIT_CONFIG_DIR"] = streamlit_config

# (Optional) also move Matplotlib & Torch caches into /tmp
os.environ["MPLCONFIGDIR"]   = os.path.join(os.environ["HOME"], ".matplotlib")
os.environ["XDG_CACHE_HOME"] = os.path.join(os.environ["HOME"], ".cache")

# 2) Prepare your own output directory under /tmp
OUTPUT_DIR = os.path.join(os.environ["HOME"], "streamlit_d2_output")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ────────────────────────────────────────────────────────────
# 3) IMPORT rest of your dependencies
# ────────────────────────────────────────────────────────────
import streamlit as st
import torch
import torchvision.transforms as T
import numpy as np
import cv2
from PIL import Image, UnidentifiedImageError
import subprocess  # For Detectron2 installation check
import sys         # For Detectron2 installation check

# Detectron2 
d2_imported_successfully = False
try:
    import detectron2
    from detectron2.engine import DefaultPredictor
    from detectron2.config import get_cfg
    from detectron2 import model_zoo
    from detectron2.utils.visualizer import Visualizer, ColorMode
    from detectron2.data import MetadataCatalog
    from detectron2.structures import Boxes  # For Bounding Boxes
    d2_imported_successfully = True
    print("Detectron2 utilities imported successfully.")
except ImportError:
    st.error("Detectron2 not found or not installed correctly. Please ensure it's installed in your environment.")
    print("Failed to import Detectron2 utilities.")
except Exception as e:
    st.error(f"An error occurred during Detectron2 imports: {e}")
    print(f"An error occurred during Detectron2 imports: {e}")

# PyTorch
from torchvision import models as torchvision_models
import torch.nn as nn

# Configuration
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
CNN_INPUT_SIZE = 224
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]
MODEL_PATH = "model/pix3d_dimension_estimator_mask_crop.pth"

# Dimension Estimation CNN
def create_dimension_estimator_cnn_for_inference(num_outputs=4):
    model = torchvision_models.resnet50(weights=None)
    num_ftrs = model.fc.in_features
    model.fc = nn.Sequential(
        nn.Linear(num_ftrs, 512),
        nn.ReLU(),
        nn.Dropout(0.5),
        nn.Linear(512, num_outputs)
    )
    return model

@st.cache_resource
def load_dimension_model():
    if not os.path.exists(MODEL_PATH):
        st.error(f"Dimension estimation model not found at {MODEL_PATH}. Please check the path.")
        return None
    try:
        model = create_dimension_estimator_cnn_for_inference()
        model.load_state_dict(torch.load(MODEL_PATH, map_location=DEVICE))
        model.to(DEVICE)
        model.eval()
        print(f"Dimension estimation model loaded from {MODEL_PATH}")
        return model
    except Exception as e:
        st.error(f"Error loading dimension estimation model: {e}")
        return None

@st.cache_resource
def load_detectron2_model():
    if not d2_imported_successfully:
        return None, None
    try:
        cfg = get_cfg()
        cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
        cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
        cfg.MODEL.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
        predictor = DefaultPredictor(cfg)
        print("Detectron2 predictor created.")
        return predictor, cfg
    except Exception as e:
        st.error(f"Error loading Detectron2 model: {e}")
        return None, None

def get_largest_instance_index(instances):
    if not len(instances):
        return -1
    if instances.has("pred_masks"):
        areas = instances.pred_masks.sum(dim=(1,2))
        return int(areas.argmax()) if len(areas) > 0 else 0
    elif instances.has("pred_boxes"):
        boxes = instances.pred_boxes.tensor
        areas = (boxes[:,2]-boxes[:,0]) * (boxes[:,3]-boxes[:,1])
        return int(areas.argmax()) if len(areas) > 0 else 0
    return 0

def crop_from_mask(image_np_rgb, mask_tensor):
    mask_np = mask_tensor.cpu().numpy().astype(np.uint8)
    if mask_np.sum() == 0: return None
    rows = np.any(mask_np, axis=1)
    cols = np.any(mask_np, axis=0)
    if not rows.any() or not cols.any(): return None
    ymin, ymax = np.where(rows)[0][[0, -1]]
    xmin, xmax = np.where(cols)[0][[0, -1]]
    pad = 5
    ymin = max(0, ymin-pad); xmin = max(0, xmin-pad)
    ymax = min(image_np_rgb.shape[0]-1, ymax+pad)
    xmax = min(image_np_rgb.shape[1]-1, xmax+pad)
    if ymin>=ymax or xmin>=xmax: return None
    return image_np_rgb[ymin:ymax+1, xmin:xmax+1]

def predict_dimensions_cnn(img_rgb, model):
    if model is None:
        return {"Length": "N/A", "Width": "N/A", "Height": "N/A", "Volume": "N/A"}
    try:
        transform = T.Compose([
            T.ToPILImage(),
            T.Resize((CNN_INPUT_SIZE, CNN_INPUT_SIZE)),
            T.ToTensor(),
            T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)
        ])
        inp = transform(img_rgb).unsqueeze(0).to(DEVICE)
        with torch.no_grad():
            out = model(inp).squeeze().cpu().tolist()
        while len(out)<4: out.append(0.0)
        L, W, H, V = out
        return {
            "Length (cm)": f"{L*100:.1f}",
            "Width (cm)" : f"{W*100:.1f}",
            "Height (cm)": f"{H*100:.1f}",
            "Volume (cm³)": f"{V*1e6:.1f}"
        }
    except Exception as e:
        print(f"CNN predict error: {e}")
        return {"Length": "Error", "Width":"Error", "Height":"Error", "Volume":"Error"}

# ────────────────────────────────────────────────────────────
# Streamlit UI
# ────────────────────────────────────────────────────────────
st.set_page_config(layout="wide", page_title="Object Dimension Estimator")
st.title("Object Dimension & Volume Estimation")

dim_model = load_dimension_model()
d2_predictor, d2_cfg = (None, None)
d2_metadata = None
if d2_imported_successfully:
    d2_predictor, d2_cfg = load_detectron2_model()
    if d2_cfg:
        try:
            d2_metadata = MetadataCatalog.get(d2_cfg.DATASETS.TRAIN[0])
        except:
            d2_metadata = MetadataCatalog.get("coco_2017_val")

uploaded = st.file_uploader("Upload an image", type=["jpg","jpeg","png"])
if uploaded:
    st.subheader(uploaded.name)
    try:
        img = Image.open(uploaded).convert("RGB")
        img_np = np.array(img)
        bgr = cv2.cvtColor(img_np, cv2.COLOR_RGB2BGR)
    except:
        st.error("Invalid image.")
        bgr = None

    if bgr is not None and d2_predictor and dim_model:
        with st.spinner("Processing..."):
            outs = d2_predictor(bgr); inst = outs["instances"].to("cpu")
            if len(inst)==0:
                st.warning("No objects detected.")
            else:
                viz = Visualizer(bgr[:,:,::-1], metadata=d2_metadata, scale=0.8, instance_mode=ColorMode.IMAGE_BW)
                out_vis = viz.draw_instance_predictions(inst)
                det_img = out_vis.get_image()[:,:,::-1]
                st.image(det_img, use_column_width=True)

                idx = get_largest_instance_index(inst)
                if idx>=0:
                    mask = inst[idx].pred_masks[0] if inst.has("pred_masks") else None
                    crop = crop_from_mask(img_np, mask) if mask is not None else None
                    if crop is not None:
                        st.image(crop, caption="Cropped Object", width=250)
                        dims = predict_dimensions_cnn(crop, dim_model)
                        st.json(dims)
                    else:
                        st.error("Could not crop object.")
    elif not d2_imported_successfully:
        st.error("Detectron2 not loaded.")
    else:
        st.error("Model not loaded.")

# Sidebar status
st.sidebar.markdown("---")
st.sidebar.write(f"Device: {DEVICE}")
st.sidebar.write(f"Detectron2: {'OK' if d2_predictor else 'Failed'}")
st.sidebar.write(f"Dim CNN: {'OK' if dim_model else 'Failed'}")