init

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,8 +1,8 @@
 ---
 title: Shic
-emoji: 🐠
+emoji: 📊
 colorFrom: purple
-colorTo: purple
+colorTo: green
 sdk: gradio
 sdk_version: 4.39.0
 app_file: app.py

app.py

@@ -0,0 +1,226 @@
+import gradio as gr
+import os
+import numpy as np
+import trimesh as tm
+from src.model import DinoV2
+from src.shape_model import CSE
+from PIL import Image, ImageDraw
+import torch
+from torchvision import transforms
+import matplotlib.pyplot as plt
+import hashlib
+
+def image_hash(image):
+    """Generate a hash for an image."""
+    image_bytes = image.tobytes()
+    hash_function = hashlib.sha256()
+    hash_function.update(image_bytes)
+    return hash_function.hexdigest()
+
+
+#device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cpu')
+models = {}
+for class_name in ['bear', 'horse', 'elephant']:
+    print(f'Loading model weights for {class_name}')
+    models[class_name] = {
+        'image_encoder': DinoV2(16),
+        'cse': CSE(class_name=class_name, num_basis=64, device=device)
+    }
+    models[class_name]['image_encoder'].load_state_dict(torch.load(f'./models/weights/{class_name}.pth', map_location=device))
+    models[class_name]['cse'].load_state_dict(torch.load(f'./models/weights/{class_name}_cse.pth', map_location=device))
+    models[class_name]['cse'].functional_basis = torch.load(f'./models/weights/{class_name}_lbo.pth', map_location=device)
+    
+    models[class_name]['image_encoder'] = models[class_name]['image_encoder'].to(device)
+    models[class_name]['cse'] = models[class_name]['cse'].to(device)
+    models[class_name]['cse'].functional_basis = models[class_name]['cse'].functional_basis.to(device)
+    models[class_name]['cse'].weight_matrix = models[class_name]['cse'].weight_matrix.to(device)
+
+    models[class_name]['shape_feats'] = models[class_name]['cse']().to(device)
+
+# Convert PIL image to a format your model expects (e.g., torch.Tensor)
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+cached_features = {'bear': {}, 'horse': {}, 'elephant': {}}
+
+text_description = """
+# Demo for SHIC: Shape-Image Correspondences with no Keypoint Superivision (ECCV 2024)
+Project website: https://www.robots.ox.ac.uk/~vgg/research/shic/
+
+- **Step 1:** First select a class (now it defaults to 'bear') 
+- **Step 2:** Upload an image of an animal from that class (or select one of the provided examples)
+- **Step 3:** Click on the image (somewhere over the object) to see the image-to-shape correspondences. You can keep clicking on the input image to see new correspondences.
+
+Notes: 
+- You can click and drag to rotate the 3D shape
+- Currently the demo supports bears, horses, and elephants. Other classes coming soon!
+- Make sure you have selected the correct class for your image (It works cross-class too though!)
+"""
+
+example_images_dir = './gradio_example_images/'
+example_images_names = [
+    'bear1.png', 'bear2.png', 'bear3.png', 'winnie.png',
+    'horse1.png', 'horse2.png', 'mylittlepony.png', 'ponyta.png',
+    'elephant1.png', 'elephant2.png', 'dumbo.png', 'phanpy.png'
+]
+example_images = [os.path.join(example_images_dir, img) for img in example_images_names]
+sphere_verts_ = torch.load(f'./models/weights/sphere_verts.pth', map_location=device)
+sphere_faces_ = torch.load(f'./models/weights/sphere_faces.pth', map_location=device)
+def center_crop(img):
+    """
+    Center crops an image to the target size of 224x224.
+    """
+    width, height = img.size   # Get dimensions
+    # Calculate the target size for center cropping
+    target_size = min(width, height)
+    
+    # Calculate the coordinates for center cropping
+    left = (width - target_size) // 2
+    top = (height - target_size) // 2
+    right = left + target_size
+    bottom = top + target_size
+    
+    # Perform center cropping
+    cropped_img = img.crop((left, top, right, bottom))
+    
+    return cropped_img
+
+def draw_point_on_image(image, x_, y_):
+    """Draws a red dot on a copy of the image at the specified point."""
+    # Make a copy of the image to avoid altering the original
+    image_copy = image.copy()
+    draw = ImageDraw.Draw(image_copy)
+    x, y = x_, y_  # Adjust these based on the actual structure of `point`
+    dot_radius = image.size[0] // 100
+    # Draw a red dot
+    draw.ellipse([(y-dot_radius, x-dot_radius), (y+dot_radius, x+dot_radius)], fill='red')
+    
+    return image_copy
+
+def rotate_y(vertices, angle_degrees):
+    angle_radians = np.radians(angle_degrees)
+    rotation_matrix = np.array([
+        [np.cos(angle_radians), 0, np.sin(angle_radians)],
+        [0, 1, 0],
+        [-np.sin(angle_radians), 0, np.cos(angle_radians)]
+    ])
+    
+    # Assuming vertices is a numpy array of shape (N, 3)
+    rotated_vertices = np.dot(vertices, rotation_matrix)
+    return rotated_vertices
+
+def make_final_mesh(verts, faces, similarities):
+    vert_argmax = similarities.argmax(dim=1)
+    vertex = verts[vert_argmax]
+    color=[255, 0, 0]
+
+    vertex_colors=similarities.transpose(1,0).cpu().detach().numpy()
+    # to viridis color map
+    vertex_colors = plt.cm.viridis(vertex_colors)[:, 0, :3]
+
+    num_verts_so_far = len(verts)
+
+    
+    # Create a sphere mesh
+
+    # Scale and translate the sphere to the desired location and size
+    scale_dot = 0.015  # radius of the sphere
+    translation = torch.tensor(vertex, device=device).unsqueeze(0)  # desired location
+
+    verts_sphere = sphere_verts_ * scale_dot + translation  # scale and translate vertices
+    faces_sphere = sphere_faces_ + num_verts_so_far  # faces are the same
+
+    verts_rgb_sphere = torch.tensor([color], device=device).expand(verts_sphere.shape[0], -1)[None] / 255  # [1, N, 3]
+
+
+    # verts and all sphere verts
+    # concat np arrays verts + verts_sphere.cpu().numpy() (4936,3) (2562,3)
+    all_verts = np.concatenate([verts, verts_sphere.cpu().numpy()], axis=0)
+    all_faces = np.concatenate([faces, faces_sphere.cpu().numpy()], axis=0)
+
+    all_textures = np.concatenate([vertex_colors, verts_rgb_sphere.cpu().numpy()[0]], axis=0)
+
+    return tm.Trimesh(vertices=all_verts, faces=all_faces, vertex_colors=all_textures)
+
+def process_mesh(image, class_name, x_, y_):
+    x_, y_ = x_, y_
+    h, w = image.size
+
+    x = torch.tensor(x_ * 224 / w)
+    y = torch.tensor(y_ * 224 / h)
+
+    hashed_image = image_hash(image)
+    if hashed_image in cached_features[class_name]:
+        feats = cached_features[class_name][hashed_image]
+    else:
+        
+        image_tensor = transform(image).unsqueeze(0)
+        
+        # Predict texture
+        feats = models[class_name]['image_encoder'](image_tensor.to(device))
+
+        cached_features[class_name][hashed_image] = feats
+
+    # print('feats shape', feats.shape)
+
+    sampled_feats = feats[:, :, x.long(), y.long()]
+    similarities = torch.einsum('ik, lk -> il', sampled_feats, models[class_name]['shape_feats'])
+    # normalize similarities
+    similarities = (similarities - similarities.min()) / (similarities.max() - similarities.min())
+
+    faces =  models[class_name]['cse'].shape['faces'].cpu().numpy().copy()
+    verts = models[class_name]['cse'].shape['verts'].cpu().numpy().copy()
+
+    # rotate the shape 235 
+    verts = rotate_y(verts, 145)
+
+    mesh = make_final_mesh(verts, faces, similarities)
+    # save as obj
+    mesh_path = './mesh.obj'
+    mesh.export(mesh_path)
+    
+    return mesh_path
+
+def update_output(image, class_name, evt: gr.SelectData):
+    if class_name is None:
+        class_name = 'bear'
+    # This function will be triggered when an image is clicked.
+    # evt contains the click event data, including the coordinates.
+    x_, y_ = evt.index[1], evt.index[0]
+    modified_image = draw_point_on_image(image, x_, y_)
+    mesh_path = process_mesh(image, class_name, x_, y_)
+    return modified_image, mesh_path  # Replace with the actual model path
+    
+
+with gr.Blocks() as demo:
+    # choose a class
+    gr.Markdown(text_description)
+
+
+    with gr.Row(variant="panel"):
+        with gr.Column(scale=1):
+            class_name = gr.Dropdown(choices=['bear', 'horse', 'elephant'], 
+                               label="Select a class (defaults to 'bear')")
+            input_img = gr.Image(label="Input", type="pil", width=256)
+            gr.Examples(
+                examples = example_images,
+                inputs = [input_img],
+                cache_examples=False,
+                label='Feel free to use one of our provided examples!',
+                examples_per_page=30
+            )
+        with gr.Column(scale=1):    
+            output_img = gr.Image(label="Selected Point", interactive=False, height=512)
+        with gr.Column(scale=1):
+            output = gr.Model3D(label='Pixel to Vertex Similarities', height=512)
+    
+
+
+    input_img.select(update_output, [input_img, class_name], [output_img, output])
+
+if __name__ == "__main__":
+    demo.launch(share=True)

models/weights/bear.pth

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

models/weights/bear_cse.pth

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

models/weights/bear_lbo.pth

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

models/weights/elephant.pth

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

models/weights/elephant_cse.pth

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

models/weights/elephant_lbo.pth

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

models/weights/horse.pth

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

models/weights/horse_cse.pth

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

models/weights/horse_lbo.pth

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

models/weights/sphere_faces.pth

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

models/weights/sphere_verts.pth

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