canva-render/scripts/post_processing.py

175 lines
7 KiB
Python
Raw Normal View History

2023-12-19 06:39:13 +00:00
import logging
import asyncio
import time
from collections import Counter
import cv2
import numpy as np
from azure.storage.blob.aio import BlobServiceClient, download_blob_from_url
from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorCursor, AsyncIOMotorCollection, AsyncIOMotorDatabase
AZUREBLOB_SAS_URL = "https://internblob.blob.core.windows.net/v-lixinyang/?sp=racwdli&st=2023-09-17T15:37:58Z&se=2023-12-31T23:37:58Z&spr=https&sv=2022-11-02&sr=c&sig=u%2FPbZ4fNttAPeLj0NEEpX0eIgFcjhot%2Bmy3iGd%2BCmxk%3D"
CONTAINER = "canva-render-10.19"
MONGODB_URI = "mongodb://localhost:27017/canva"
class BlobAsync(object):
async def readall(self, blob):
blob_service_client = BlobServiceClient(AZUREBLOB_SAS_URL)
async with blob_service_client:
container_client = blob_service_client.get_container_client(CONTAINER)
# async for bname in container_client.list_blob_names():
# print(bname)
blob_client = container_client.get_blob_client(blob)
if not await blob_client.exists():
return None
stream = await blob_client.download_blob()
return stream.readall()
async def open_image(self, blob: str):
async with BlobServiceClient(AZUREBLOB_SAS_URL) as blob_service_client:
container_client = blob_service_client.get_container_client(CONTAINER)
# async for bname in container_client.list_blob_names():
# print(bname)
blob_client = container_client.get_blob_client(blob)
if not await blob_client.exists():
return None
stream = await blob_client.download_blob()
buf = np.frombuffer(await stream.readall(), dtype=np.uint8)
image = cv2.imdecode(buf, cv2.IMREAD_COLOR)
await blob_client.close()
await container_client.close()
return image
async def upload_image(self, blob, image):
async with BlobServiceClient(AZUREBLOB_SAS_URL) as blob_service_client:
# Instantiate a new ContainerClient
container_client = blob_service_client.get_container_client(CONTAINER)
blob_client = container_client.get_blob_client(blob)
is_success, buffer = cv2.imencode('.png', image)
await blob_client.upload_blob(data=buffer.tobytes(), overwrite=True)
await blob_client.close()
await container_client.close()
async def get_mask(img1, img2):
"""Assume img1 and img2 are exactly the same, except text areas
"""
try:
diff = cv2.absdiff(img1, img2)
except:
raise ValueError("img1 and img2 are not the same size")
mask = cv2.cvtColor(diff, cv2.COLOR_RGBA2GRAY)
thresh, binmask= cv2.threshold(mask, 10, 255, cv2.THRESH_BINARY)
return thresh, binmask
async def filter_mask_size(mask, thresh=0.4):
non_zero_pixels = cv2.countNonZero(mask)
total_pixels = mask.shape[0] * mask.shape[1]
if non_zero_pixels > total_pixels * thresh:
return True
else:
return False
mask_filtered_count = Counter()
async def process_cdf(blob: BlobAsync, collection, cdf):
folder = cdf["rendered_folder"]
async with asyncio.TaskGroup() as g:
task1 = g.create_task(blob.open_image(f"{folder}/t=true.png"))
task2 = g.create_task(blob.open_image(f"{folder}/t=false.png"))
img1, img2 = task1.result(), task2.result()
if img1 is None and img2 is None:
mask_filtered_count["not found"] += 1
await collection.update_one({"_id": cdf["_id"]}, {"$set": {"last_mask_render": -1, "failed_reason": "not found both"}})
return
if img1 is None:
mask_filtered_count["not found"] += 1
await collection.update_one({"_id": cdf["_id"]}, {"$set": {"last_mask_render": -1, "failed_reason": "not found t=true"}})
return
if img2 is None:
mask_filtered_count["not found"] += 1
await collection.update_one({"_id": cdf["_id"]}, {"$set": {"last_mask_render": -1, "failed_reason": "not found t=false"}})
return
try:
binary_thresh, mask = await get_mask(img1, img2)
except ValueError as e:
await collection.update_one({"_id": cdf["_id"]}, {"$set": {"last_mask_render": -1, "failed_reason": "size not match"}})
mask_filtered_count["size not match"] += 1
return
mask_filters = [
(filter_mask_size, "mask too small")
]
tasks = list()
async with asyncio.TaskGroup() as g:
for f, reason in mask_filters:
tasks.append((g.create_task(f(mask)), reason))
for task, reason in tasks:
if task.result():
mask_filtered_count[reason] += 1
await collection.update_one({"_id": cdf["_id"]}, {"$set": {"last_mask_render": -1, "failed_reason": reason}, "$unset": {"last_fetched": -1}})
return
await blob.upload_image(f"{folder}/mask.png", mask)
await collection.update_one({"_id": cdf["_id"]}, {"$set": {"last_mask_render": time.time()}})
mask_filtered_count["success"] += 1
async def main():
client = AsyncIOMotorClient(MONGODB_URI)
db = client.get_database("canva")
collection = db["cdf"]
logger = logging.getLogger('azure.mgmt.resource')
logger.setLevel(logging.WARNING)
blob = BlobAsync()
cdf_cursor: AsyncIOMotorCursor = collection.find({
'$or': [
{ '$and': [
{ 'rendered_folder': { '$exists': True } },
{ 'last_fetched': { '$gt': 1697688216 } },
{ 'last_fetched': { '$lt': time.time() - 600 } },
{ 'last_mask_render': { '$exists': False }}
]},
{ '$and': [
{ 'last_fetched': {'$gt': 1697998932}},
{ 'last_mask_render': { '$not': { '$gt': 0 } } }
]}
]}, batch_size=400)
cdf_list = await cdf_cursor.to_list(length=200)
await cdf_cursor.close()
while cdf_list is not []:
async with asyncio.TaskGroup() as g:
taskset = set()
for cdf in cdf_list:
taskset.add(
g.create_task(process_cdf(blob, collection, cdf))
)
await asyncio.sleep(10)
cdf_cursor: AsyncIOMotorCursor = collection.find({
'$or': [
{ '$and': [
{ 'rendered_folder': { '$exists': True } },
{ 'last_fetched': { '$gt': 1697688216 } },
{ 'last_fetched': { '$lt': time.time() - 600 } },
{ 'last_mask_render': { '$exists': False }}
]},
{ '$and': [
{ 'last_fetched': {'$gt': 1697998932}},
{ 'last_mask_render': { '$not': { '$gt': 0 } } }
]}
]}, batch_size=400)
cdf_list = await cdf_cursor.to_list(length=200)
await cdf_cursor.close()
print(mask_filtered_count)
if __name__ == "__main__":
asyncio.run(main())