canadianjosieharrison commited on
Commit
baffb91
1 Parent(s): 73351fb

Upload 11 files

Browse files
.gitattributes CHANGED
@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ image/render_25.jpg filter=lfs diff=lfs merge=lfs -text
37
+ image/render1.png filter=lfs diff=lfs merge=lfs -text
app.py ADDED
@@ -0,0 +1,94 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import base64
2
+ import io
3
+ from fastapi import FastAPI, UploadFile, File, HTTPException
4
+ import os
5
+ import shutil
6
+ from PIL import Image
7
+ from fastapi.responses import JSONResponse
8
+ from semantic_seg_model import segmentation_inference
9
+ from similarity_inference import similarity_inference
10
+ import json
11
+ from gradio_client import Client, file
12
+
13
+ app = FastAPI()
14
+
15
+ ## Initialize the pipeline
16
+ input_images_dir = "image/"
17
+ temp_processing_dir = input_images_dir + "processed/"
18
+
19
+ # Define a function to handle the POST request at `imageAnalyzer`
20
+ @app.post("/imageAnalyzer")
21
+ def imageAnalyzer(image: UploadFile = File(...)):
22
+ """
23
+ This function takes in an image filepath and will return the PolyHaven url addresses of the
24
+ top k materials similar to the wall, ceiling, and floor.
25
+ """
26
+ try:
27
+ # load image
28
+ image_path = os.path.join(input_images_dir, image.filename)
29
+ with open(image_path, "wb") as buffer:
30
+ shutil.copyfileobj(image.file, buffer)
31
+ image = Image.open(image_path)
32
+
33
+ # segment into components
34
+ segmentation_inference(image, temp_processing_dir)
35
+
36
+ # identify similar materials for each component
37
+ matching_urls = similarity_inference(temp_processing_dir)
38
+ print(matching_urls)
39
+
40
+ # Return the urls in a JSON response
41
+ return matching_urls
42
+
43
+ except Exception as e:
44
+ print(str(e))
45
+ raise HTTPException(status_code=500, detail=str(e))
46
+
47
+
48
+ client = Client("MykolaL/StableDesign")
49
+
50
+ @app.post("/image-render")
51
+ def imageRender(prompt: str, image: UploadFile = File(...)):
52
+ """
53
+ Makes a prediction using the "StableDesign" model hosted on a server.
54
+
55
+ Returns:
56
+ The prediction result.
57
+ """
58
+ try:
59
+ image_path = os.path.join(input_images_dir, image.filename)
60
+ with open(image_path, "wb") as buffer:
61
+ shutil.copyfileobj(image.file, buffer)
62
+ image = Image.open(image_path)
63
+ # Convert PIL image to the required format for the prediction model, if necessary
64
+ # This example assumes the model accepts PIL images directly
65
+
66
+ result = client.predict(
67
+ image=file(image_path),
68
+ text=prompt,
69
+ num_steps=50,
70
+ guidance_scale=10,
71
+ seed=1111664444,
72
+ strength=0.9,
73
+ a_prompt="interior design, 4K, high resolution, photorealistic",
74
+ n_prompt="window, door, low resolution, banner, logo, watermark, text, deformed, blurry, out of focus, surreal, ugly, beginner",
75
+ img_size=768,
76
+ api_name="/on_submit"
77
+ )
78
+ image_path = result
79
+ if not os.path.exists(image_path):
80
+ raise HTTPException(status_code=404, detail="Image not found")
81
+
82
+ # Open the image file and convert it to base64
83
+ with open(image_path, "rb") as img_file:
84
+ base64_str = base64.b64encode(img_file.read()).decode('utf-8')
85
+
86
+ return JSONResponse(content={"image": base64_str}, status_code=200)
87
+ except Exception as e:
88
+ print(str(e))
89
+ raise HTTPException(status_code=500, detail=str(e))
90
+
91
+
92
+ @app.get("/")
93
+ def test():
94
+ return {"Hello": "World"}
data/embedding_db.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:014a29f36b6fbd851650ef9dd1e13c4fd2a9bcdb5d5776b920a343ed787c3de7
3
+ size 3017909
image/processed/ceiling2.png ADDED
image/processed/floor1.png ADDED
image/processed/wall0.png ADDED
image/render1.png ADDED

Git LFS Details

  • SHA256: b351f0e7ea038cbdb42f43f8cb02eb2b063ce44ccb3c24be42ebf41358b8042c
  • Pointer size: 132 Bytes
  • Size of remote file: 2.1 MB
image/render_25.jpg ADDED

Git LFS Details

  • SHA256: 6096e70a8c3fbf10cdf55e1b89b42e9e65912c7469d05659ebec0330aecd588a
  • Pointer size: 132 Bytes
  • Size of remote file: 2.9 MB
image_helpers.py ADDED
@@ -0,0 +1,113 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from PIL import Image
3
+ from cv2 import imread, cvtColor, COLOR_BGR2GRAY, COLOR_BGR2BGRA, COLOR_BGRA2RGB, threshold, THRESH_BINARY_INV, findContours, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE, contourArea, minEnclosingCircle
4
+ import numpy as np
5
+ import torch
6
+ import matplotlib.pyplot as plt
7
+
8
+ def convert_images_to_grayscale(folder_path):
9
+ # Check if the folder exists
10
+ if not os.path.isdir(folder_path):
11
+ print(f"The folder path {folder_path} does not exist.")
12
+ return
13
+
14
+ # Iterate over all files in the folder
15
+ for filename in os.listdir(folder_path):
16
+ if filename.endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif')):
17
+ image_path = os.path.join(folder_path, filename)
18
+
19
+ # Open an image file
20
+ with Image.open(image_path) as img:
21
+ # Convert image to grayscale
22
+ grayscale_img = img.convert('L').convert('RGB')
23
+ grayscale_img.save(os.path.join(folder_path, filename))
24
+
25
+ def crop_center_largest_contour(folder_path):
26
+ for each_image in os.listdir(folder_path):
27
+ image_path = os.path.join(folder_path, each_image)
28
+ image = imread(image_path)
29
+ gray_image = cvtColor(image, COLOR_BGR2GRAY)
30
+
31
+ # Threshold the image to get the non-white pixels
32
+ _, binary_mask = threshold(gray_image, 254, 255, THRESH_BINARY_INV)
33
+
34
+ # Find the largest contour
35
+ contours, _ = findContours(binary_mask, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE)
36
+ largest_contour = max(contours, key=contourArea)
37
+
38
+ # Get the minimum enclosing circle
39
+ (x, y), radius = minEnclosingCircle(largest_contour)
40
+ center = (int(x), int(y))
41
+ radius = int(radius/3) # Divide by three (arbitrary) to make shape better
42
+
43
+ # Crop the image to the bounding box of the circle
44
+ x_min = max(0, center[0] - radius)
45
+ x_max = min(image.shape[1], center[0] + radius)
46
+ y_min = max(0, center[1] - radius)
47
+ y_max = min(image.shape[0], center[1] + radius)
48
+ cropped_image = image[y_min:y_max, x_min:x_max]
49
+ cropped_image_rgba = cvtColor(cropped_image, COLOR_BGR2BGRA)
50
+ cropped_pil_image = Image.fromarray(cvtColor(cropped_image_rgba, COLOR_BGRA2RGB))
51
+ cropped_pil_image.save(image_path)
52
+
53
+ def extract_embeddings(transformation_chain, model: torch.nn.Module):
54
+ """Utility to compute embeddings."""
55
+ device = model.device
56
+
57
+ def pp(batch):
58
+ images = batch["image"]
59
+ image_batch_transformed = torch.stack(
60
+ [transformation_chain(image) for image in images]
61
+ )
62
+ new_batch = {"pixel_values": image_batch_transformed.to(device)}
63
+ with torch.no_grad():
64
+ embeddings = model(**new_batch).last_hidden_state[:, 0].cpu()
65
+ return {"embeddings": embeddings}
66
+
67
+ return pp
68
+
69
+ def compute_scores(emb_one, emb_two):
70
+ """Computes cosine similarity between two vectors."""
71
+ scores = torch.nn.functional.cosine_similarity(emb_one, emb_two)
72
+ return scores.numpy().tolist()
73
+
74
+
75
+ def fetch_similar(image, transformation_chain, device, model, all_candidate_embeddings, candidate_ids, top_k=3):
76
+ """Fetches the `top_k` similar images with `image` as the query."""
77
+ # Prepare the input query image for embedding computation.
78
+ image_transformed = transformation_chain(image).unsqueeze(0)
79
+ new_batch = {"pixel_values": image_transformed.to(device)}
80
+
81
+ # Compute the embedding.
82
+ with torch.no_grad():
83
+ query_embeddings = model(**new_batch).last_hidden_state[:, 0].cpu()
84
+
85
+ # Compute similarity scores with all the candidate images at one go.
86
+ # We also create a mapping between the candidate image identifiers
87
+ # and their similarity scores with the query image.
88
+ sim_scores = compute_scores(all_candidate_embeddings, query_embeddings)
89
+ similarity_mapping = dict(zip(candidate_ids, sim_scores))
90
+
91
+ # Sort the mapping dictionary and return `top_k` candidates.
92
+ similarity_mapping_sorted = dict(
93
+ sorted(similarity_mapping.items(), key=lambda x: x[1], reverse=True)
94
+ )
95
+ id_entries = list(similarity_mapping_sorted.keys())[:top_k]
96
+
97
+ ids = list(map(lambda x: int(x.split("_")[0]), id_entries))
98
+ return ids
99
+
100
+ def plot_images(images):
101
+
102
+ plt.figure(figsize=(20, 10))
103
+ columns = 6
104
+ for (i, image) in enumerate(images):
105
+ ax = plt.subplot(int(len(images) / columns + 1), columns, i + 1)
106
+ if i == 0:
107
+ ax.set_title("Query Image\n")
108
+ else:
109
+ ax.set_title(
110
+ "Similar Image # " + str(i)
111
+ )
112
+ plt.imshow(np.array(image).astype("int"))
113
+ plt.axis("off")
requirements.txt ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ aiohttp==3.9.5
2
+ aiosignal==1.3.1
3
+ annotated-types==0.7.0
4
+ anyio==4.4.0
5
+ asttokens==2.4.1
6
+ attrs==23.2.0
7
+ certifi==2024.2.2
8
+ charset-normalizer==3.3.2
9
+ click==8.1.7
10
+ colorama==0.4.6
11
+ comm==0.2.2
12
+ contourpy==1.2.1
13
+ cycler==0.12.1
14
+ datasets==2.19.1
15
+ debugpy==1.8.1
16
+ decorator==5.1.1
17
+ dill==0.3.8
18
+ dnspython==2.6.1
19
+ email_validator==2.1.1
20
+ executing==2.0.1
21
+ fastapi==0.111.0
22
+ fastapi-cli==0.0.4
23
+ filelock==3.14.0
24
+ fonttools==4.52.1
25
+ frozenlist==1.4.1
26
+ fsspec==2024.3.1
27
+ gradio_client==0.17.0
28
+ h11==0.14.0
29
+ httpcore==1.0.5
30
+ httptools==0.6.1
31
+ httpx==0.27.0
32
+ huggingface-hub==0.23.1
33
+ idna==3.7
34
+ intel-openmp==2021.4.0
35
+ ipykernel==6.29.4
36
+ ipython==8.24.0
37
+ jedi==0.19.1
38
+ Jinja2==3.1.4
39
+ jupyter_client==8.6.2
40
+ jupyter_core==5.7.2
41
+ kiwisolver==1.4.5
42
+ markdown-it-py==3.0.0
43
+ MarkupSafe==2.1.5
44
+ matplotlib==3.9.0
45
+ matplotlib-inline==0.1.7
46
+ mdurl==0.1.2
47
+ mkl==2021.4.0
48
+ mpmath==1.3.0
49
+ multidict==6.0.5
50
+ multiprocess==0.70.16
51
+ nest-asyncio==1.6.0
52
+ networkx==3.3
53
+ numpy==1.26.4
54
+ opencv-python-headless==4.9.0.80
55
+ orjson==3.10.3
56
+ packaging==24.0
57
+ pandas==2.2.2
58
+ parso==0.8.4
59
+ pillow==10.3.0
60
+ platformdirs==4.2.2
61
+ prompt-toolkit==3.0.43
62
+ psutil==5.9.8
63
+ pure-eval==0.2.2
64
+ pyarrow==16.1.0
65
+ pyarrow-hotfix==0.6
66
+ pydantic==2.7.1
67
+ pydantic_core==2.18.2
68
+ Pygments==2.18.0
69
+ pyparsing==3.1.2
70
+ python-dateutil==2.9.0.post0
71
+ python-dotenv==1.0.1
72
+ python-multipart==0.0.9
73
+ # pytz==2024.1
74
+ # pywin32==306
75
+ PyYAML==6.0.1
76
+ pyzmq==26.0.3
77
+ regex==2024.5.15
78
+ requests==2.32.2
79
+ rich==13.7.1
80
+ safetensors==0.4.3
81
+ shellingham==1.5.4
82
+ six==1.16.0
83
+ sniffio==1.3.1
84
+ stack-data==0.6.3
85
+ starlette==0.37.2
86
+ sympy==1.12
87
+ tbb==2021.12.0
88
+ tokenizers==0.19.1
89
+ torch==2.3.0
90
+ torchvision==0.18.0
91
+ tornado==6.4
92
+ tqdm==4.66.4
93
+ traitlets==5.14.3
94
+ transformers==4.41.1
95
+ typer==0.12.3
96
+ typing_extensions==4.12.0
97
+ tzdata==2024.1
98
+ ujson==5.10.0
99
+ urllib3==2.2.1
100
+ uvicorn==0.29.0
101
+ watchfiles==0.21.0
102
+ # wcwidth==0.2.13
103
+ # websockets==12.0
104
+ xxhash==3.4.1
105
+ yarl==1.9.4
semantic_seg_model.py ADDED
@@ -0,0 +1,317 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import pipeline, AutoImageProcessor, SegformerForSemanticSegmentation
3
+ from typing import List
4
+ from PIL import Image, ImageDraw, ImageFont, ImageChops, ImageMorph
5
+ import numpy as np
6
+ import datasets
7
+
8
+ def find_center_of_non_black_pixels(image):
9
+ # Get image dimensions
10
+ width, height = image.size
11
+
12
+ # Iterate over the pixels to find the center of the non-black pixels
13
+ total_x = 0
14
+ total_y = 0
15
+ num_non_black_pixels = 0
16
+ top, left, bottom, right = height, width, 0, 0
17
+ for y in range(height):
18
+ for x in range(width):
19
+ pixel = image.getpixel((x, y))
20
+ if pixel != (255, 255, 255): # Non-black pixel
21
+ total_x += x
22
+ total_y += y
23
+ num_non_black_pixels += 1
24
+ top = min(top, y)
25
+ left = min(left, x)
26
+ bottom = max(bottom, y)
27
+ right = max(right, x)
28
+
29
+ bbox_width = right - left
30
+ bbox_height = bottom - top
31
+ bbox_size = max(bbox_height, bbox_width)
32
+ # Calculate the center of the non-black pixels
33
+ if num_non_black_pixels == 0:
34
+ return None # No non-black pixels found
35
+ center_x = total_x // num_non_black_pixels
36
+ center_y = total_y // num_non_black_pixels
37
+ return (center_x, center_y), bbox_size
38
+
39
+ def create_centered_image(image, center, bbox_size):
40
+ # Get image dimensions
41
+ width, height = image.size
42
+
43
+ # Calculate the offset to center the non-black pixels in the new image
44
+ offset_x = bbox_size // 2 - center[0]
45
+ offset_y = bbox_size // 2 - center[1]
46
+
47
+ # Create a new image with the same size as the original image
48
+ new_image = Image.new("RGB", (bbox_size, bbox_size), color=(255, 255, 255))
49
+
50
+ # Paste the non-black pixels onto the new image
51
+ new_image.paste(image, (offset_x, offset_y))
52
+
53
+ return new_image
54
+
55
+ def ade_palette():
56
+ """ADE20K palette that maps each class to RGB values."""
57
+ return [
58
+ [180, 120, 20],
59
+ [180, 120, 120],
60
+ [6, 230, 230],
61
+ [80, 50, 50],
62
+ [4, 200, 3],
63
+ [120, 120, 80],
64
+ [140, 140, 140],
65
+ [204, 5, 255],
66
+ [230, 230, 230],
67
+ [4, 250, 7],
68
+ [224, 5, 255],
69
+ [235, 255, 7],
70
+ [150, 5, 61],
71
+ [120, 120, 70],
72
+ [8, 255, 51],
73
+ [255, 6, 82],
74
+ [143, 255, 140],
75
+ [204, 255, 4],
76
+ [255, 51, 7],
77
+ [204, 70, 3],
78
+ [0, 102, 200],
79
+ [61, 230, 250],
80
+ [255, 6, 51],
81
+ [11, 102, 255],
82
+ [255, 7, 71],
83
+ [255, 9, 224],
84
+ [9, 7, 230],
85
+ [220, 220, 220],
86
+ [255, 9, 92],
87
+ [112, 9, 255],
88
+ [8, 255, 214],
89
+ [7, 255, 224],
90
+ [255, 184, 6],
91
+ [10, 255, 71],
92
+ [255, 41, 10],
93
+ [7, 255, 255],
94
+ [224, 255, 8],
95
+ [102, 8, 255],
96
+ [255, 61, 6],
97
+ [255, 194, 7],
98
+ [255, 122, 8],
99
+ [0, 255, 20],
100
+ [255, 8, 41],
101
+ [255, 5, 153],
102
+ [6, 51, 255],
103
+ [235, 12, 255],
104
+ [160, 150, 20],
105
+ [0, 163, 255],
106
+ [140, 140, 140],
107
+ [250, 10, 15],
108
+ [20, 255, 0],
109
+ [31, 255, 0],
110
+ [255, 31, 0],
111
+ [255, 224, 0],
112
+ [153, 255, 0],
113
+ [0, 0, 255],
114
+ [255, 71, 0],
115
+ [0, 235, 255],
116
+ [0, 173, 255],
117
+ [31, 0, 255],
118
+ [11, 200, 200],
119
+ [255, 82, 0],
120
+ [0, 255, 245],
121
+ [0, 61, 255],
122
+ [0, 255, 112],
123
+ [0, 255, 133],
124
+ [255, 0, 0],
125
+ [255, 163, 0],
126
+ [255, 102, 0],
127
+ [194, 255, 0],
128
+ [0, 143, 255],
129
+ [51, 255, 0],
130
+ [0, 82, 255],
131
+ [0, 255, 41],
132
+ [0, 255, 173],
133
+ [10, 0, 255],
134
+ [173, 255, 0],
135
+ [0, 255, 153],
136
+ [255, 92, 0],
137
+ [255, 0, 255],
138
+ [255, 0, 245],
139
+ [255, 0, 102],
140
+ [255, 173, 0],
141
+ [255, 0, 20],
142
+ [255, 184, 184],
143
+ [0, 31, 255],
144
+ [0, 255, 61],
145
+ [0, 71, 255],
146
+ [255, 0, 204],
147
+ [0, 255, 194],
148
+ [0, 255, 82],
149
+ [0, 10, 255],
150
+ [0, 112, 255],
151
+ [51, 0, 255],
152
+ [0, 194, 255],
153
+ [0, 122, 255],
154
+ [0, 255, 163],
155
+ [255, 153, 0],
156
+ [0, 255, 10],
157
+ [255, 112, 0],
158
+ [143, 255, 0],
159
+ [82, 0, 255],
160
+ [163, 255, 0],
161
+ [255, 235, 0],
162
+ [8, 184, 170],
163
+ [133, 0, 255],
164
+ [0, 255, 92],
165
+ [184, 0, 255],
166
+ [255, 0, 31],
167
+ [0, 184, 255],
168
+ [0, 214, 255],
169
+ [255, 0, 112],
170
+ [92, 255, 0],
171
+ [0, 224, 255],
172
+ [112, 224, 255],
173
+ [70, 184, 160],
174
+ [163, 0, 255],
175
+ [153, 0, 255],
176
+ [71, 255, 0],
177
+ [255, 0, 163],
178
+ [255, 204, 0],
179
+ [255, 0, 143],
180
+ [0, 255, 235],
181
+ [133, 255, 0],
182
+ [255, 0, 235],
183
+ [245, 0, 255],
184
+ [255, 0, 122],
185
+ [255, 245, 0],
186
+ [10, 190, 212],
187
+ [214, 255, 0],
188
+ [0, 204, 255],
189
+ [20, 0, 255],
190
+ [255, 255, 0],
191
+ [0, 153, 255],
192
+ [0, 41, 255],
193
+ [0, 255, 204],
194
+ [41, 0, 255],
195
+ [41, 255, 0],
196
+ [173, 0, 255],
197
+ [0, 245, 255],
198
+ [71, 0, 255],
199
+ [122, 0, 255],
200
+ [0, 255, 184],
201
+ [0, 92, 255],
202
+ [184, 255, 0],
203
+ [0, 133, 255],
204
+ [255, 214, 0],
205
+ [25, 194, 194],
206
+ [102, 255, 0],
207
+ [92, 0, 255],
208
+ ]
209
+
210
+ def label_to_color_image(label, colormap):
211
+ if label.ndim != 2:
212
+ raise ValueError("Expect 2-D input label")
213
+
214
+ if np.max(label) >= len(colormap):
215
+ raise ValueError("label value too large.")
216
+
217
+ return colormap[label]
218
+
219
+ labels_list = []
220
+
221
+ with open(r'labels.txt', 'r') as fp:
222
+ for line in fp:
223
+ labels_list.append(line[:-1])
224
+
225
+ colormap = np.asarray(ade_palette())
226
+ LABEL_NAMES = np.asarray(labels_list)
227
+ LABEL_TO_INDEX = {label: i for i, label in enumerate(labels_list)}
228
+ FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
229
+ FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP, colormap)
230
+ FONT = ImageFont.truetype("Arial.ttf", 1000)
231
+
232
+ def lift_black_value(image, lift_amount):
233
+ """
234
+ Increase the black values of an image by a specified amount.
235
+
236
+ Parameters:
237
+ image (PIL.Image): The image to adjust.
238
+ lift_amount (int): The amount to increase the brightness of the darker pixels.
239
+
240
+ Returns:
241
+ PIL.Image: The adjusted image with lifted black values.
242
+ """
243
+ # Ensure that we don't go out of the 0-255 range for any pixel value
244
+ def adjust_value(value):
245
+ return min(255, max(0, value + lift_amount))
246
+
247
+ # Apply the point function to each channel
248
+ return image.point(adjust_value)
249
+
250
+ torch.set_grad_enabled(False)
251
+
252
+ DEVICE = 'cuda' if torch.cuda.is_available() else "cpu"
253
+ # MIN_AREA_THRESHOLD = 0.01
254
+
255
+ pipe = pipeline("image-segmentation", model="nvidia/segformer-b5-finetuned-ade-640-640")
256
+
257
+ def segmentation_inference(
258
+ image_rgb_pil: Image.Image,
259
+ savepath: str
260
+ ):
261
+ outputs = pipe(image_rgb_pil, points_per_batch=32)
262
+
263
+ for i, prediction in enumerate(outputs):
264
+ label = prediction['label']
265
+ if (label == "floor") | (label == "wall") | (label == "ceiling"):
266
+ mask = prediction['mask']
267
+
268
+ ## Save mask
269
+ label_savepath = savepath + label + str(i) + '.png'
270
+ fill_image = Image.new("RGB", image_rgb_pil.size, color=(255,255,255))
271
+ cutout_image = Image.composite(image_rgb_pil, fill_image, mask)
272
+
273
+ # Crop mask
274
+ center, bbox_size = find_center_of_non_black_pixels(cutout_image)
275
+ if center is not None:
276
+ centered_image = create_centered_image(cutout_image, center, bbox_size)
277
+ centered_image.save(label_savepath)
278
+
279
+ ## Inspect masks
280
+ # inverted_mask = ImageChops.invert(mask)
281
+ # mask_adjusted = lift_black_value(inverted_mask, 100)
282
+ # color_index = LABEL_TO_INDEX[label]
283
+ # color = tuple(FULL_COLOR_MAP[color_index][0])
284
+ # fill_image = Image.new("RGB", image_rgb_pil.size, color=color)
285
+ # image_rgb_pil = Image.composite(image_rgb_pil, fill_image, mask_adjusted)
286
+
287
+ # Display the final image
288
+ # image_rgb_pil.show()
289
+
290
+ def online_segmentation_inference(
291
+ image_rgb_pil: Image.Image
292
+ ):
293
+ outputs = pipe(image_rgb_pil, points_per_batch=32)
294
+
295
+ # Create an image dictionary
296
+ image_dict = {"image": [], "label":[]}
297
+
298
+ for i, prediction in enumerate(outputs):
299
+ label = prediction['label']
300
+ if (label == "floor") | (label == "wall") | (label == "ceiling"):
301
+ mask = prediction['mask']
302
+
303
+ fill_image = Image.new("RGB", image_rgb_pil.size, color=(255,255,255))
304
+ cutout_image = Image.composite(image_rgb_pil, fill_image, mask)
305
+
306
+ # Crop mask
307
+ center, bbox_size = find_center_of_non_black_pixels(cutout_image)
308
+ if center is not None:
309
+ centered_image = create_centered_image(cutout_image, center, bbox_size)
310
+
311
+ # Add image to image dictionary
312
+ image_dict["image"].append(centered_image)
313
+ image_dict["label"].append(label)
314
+
315
+ segmented_ds = datasets.Dataset.from_dict(image_dict).cast_column("image", datasets.Image())
316
+ return segmented_ds
317
+
similarity_inference.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from image_helpers import convert_images_to_grayscale, crop_center_largest_contour, fetch_similar
2
+ import datasets as ds
3
+ import re
4
+ import torchvision.transforms as T
5
+ from transformers import AutoModel, AutoFeatureExtractor
6
+ import torch
7
+ import random
8
+
9
+ def similarity_inference(directory):
10
+ convert_images_to_grayscale(directory)
11
+ crop_center_largest_contour(directory)
12
+
13
+ # define processing variables needed for embedding calculation
14
+ root_directory = "data/" #"C:/Users/josie/OneDrive - Chalmers/Documents/Speckle hackathon/data/"
15
+ model_ckpt = "nateraw/vit-base-beans" ## FIND DIFFERENT MODEL
16
+ candidate_subset_emb = ds.load_dataset("canadianjosieharrison/2024hackathonembeddingdb")['train']
17
+ extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
18
+ model = AutoModel.from_pretrained(model_ckpt)
19
+ transformation_chain = T.Compose(
20
+ [
21
+ # We first resize the input image to 256x256 and then we take center crop.
22
+ T.Resize(int((256 / 224) * extractor.size["height"])),
23
+ T.CenterCrop(extractor.size["height"]),
24
+ T.ToTensor(),
25
+ T.Normalize(mean=extractor.image_mean, std=extractor.image_std),
26
+ ])
27
+ device = "cuda" if torch.cuda.is_available() else "cpu"
28
+ pt_directory = root_directory + "embedding_db.pt" #"materials/embedding_db.pt"
29
+ all_candidate_embeddings = torch.load(pt_directory, map_location=device, weights_only=True)
30
+ candidate_ids = []
31
+ for id in range(len(candidate_subset_emb)):
32
+ # Create a unique indentifier.
33
+ entry = str(id) + "_" + str(random.random()).split('.')[1]
34
+ candidate_ids.append(entry)
35
+
36
+ # load all components
37
+ test_ds = ds.load_dataset("imagefolder", data_dir=directory)
38
+ label_filenames = ds.load_dataset("imagefolder", data_dir=directory).cast_column("image", ds.Image(decode=False))
39
+
40
+ # loop through each component and return top 3 most similar
41
+ match_dict = {"ceiling": [], "floor": [], "wall": []}
42
+ for i, each_component in enumerate(test_ds['train']):
43
+ query_image = each_component["image"]
44
+ component_label = label_filenames['train'][i]['image']['path'].split('_')[-1]
45
+ print(component_label)
46
+ match = re.search(r"([a-zA-Z]+)\d*\.png", component_label)
47
+ component_label = match.group(1)
48
+ sim_ids = fetch_similar(query_image, transformation_chain, device, model, all_candidate_embeddings, candidate_ids)
49
+ for each_match in sim_ids:
50
+ texture_filename = candidate_subset_emb[each_match]['filenames']
51
+ image_url = f'https://cdn.polyhaven.com/asset_img/thumbs/{texture_filename}?width=256&height=256'
52
+ match_dict[component_label].append(image_url)
53
+
54
+ return match_dict