Spaces:
Runtime error
Runtime error
canadianjosieharrison
commited on
Commit
•
baffb91
1
Parent(s):
73351fb
Upload 11 files
Browse files- .gitattributes +2 -0
- app.py +94 -0
- data/embedding_db.pt +3 -0
- image/processed/ceiling2.png +0 -0
- image/processed/floor1.png +0 -0
- image/processed/wall0.png +0 -0
- image/render1.png +3 -0
- image/render_25.jpg +3 -0
- image_helpers.py +113 -0
- requirements.txt +105 -0
- semantic_seg_model.py +317 -0
- similarity_inference.py +54 -0
.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
|
image/render_25.jpg
ADDED
Git LFS Details
|
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
|